1
0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-10-24 02:49:04 +03:00

feat(assets): split asset records from content (#16295)

* review-stack 1/4: code (37 files, +3217/-3958)

Review-and-land stack for synap5e/feat/asset-record-content-split, generated by review-stack.py. Once approved,
merges DOWN into the layer below (a fast-forward); only the bottom layer
squash-merges into the real base. See ~/adocs/review-stack.md.
Rule: path not under tests-unit/ or tests/
Question: Is the logic change right?
Source tip: 7007d18582
Merge-base: 783545f689

* review-stack 2/4: tests-removed (24 files, +274/-8220)

Review-and-land stack for synap5e/feat/asset-record-content-split, generated by review-stack.py. Once approved,
merges DOWN into the layer below (a fast-forward); only the bottom layer
squash-merges into the real base. See ~/adocs/review-stack.md.
Rule: test file deleted, or modified with deleted/(added+deleted) >= 0.9
Question: For each dropped assertion: obsolete by a ruling, or covered by a tests-new test?
Source tip: 7007d18582
Merge-base: 783545f689

* review-stack 3/4: tests-changed (13 files, +1043/-1218)

Review-and-land stack for synap5e/feat/asset-record-content-split, generated by review-stack.py. Once approved,
merges DOWN into the layer below (a fast-forward); only the bottom layer
squash-merges into the real base. See ~/adocs/review-stack.md.
Rule: remaining modified test files (incl. conftest.py / helpers)
Question: Did the edits weaken an existing check?
Source tip: 7007d18582
Merge-base: 783545f689

* review-stack 4/4: tests-new (46 files, +8601/-0)

Review-and-land stack for synap5e/feat/asset-record-content-split, generated by review-stack.py. Once approved,
merges DOWN into the layer below (a fast-forward); only the bottom layer
squash-merges into the real base. See ~/adocs/review-stack.md.
Rule: test file added
Question: Is the code layer well covered?
Source tip: 7007d18582
Merge-base: 783545f689

* review-stack 5/6: code (13 files, +351/-104)

Review-and-land stack for synap5e/feat/assets-di, generated by review-stack.py. Once approved,
merges DOWN into the layer below (a fast-forward); only the bottom layer
squash-merges into the real base. See ~/adocs/review-stack.md.
Rule: path not under tests-unit/ or tests/
Question: Is the logic change right?
Source tip: eca2c74bff
Merge-base: 20d59d2a5f

* review-stack 6/6: tests (8 files, +753/-238)

Review-and-land stack for synap5e/feat/assets-di, generated by review-stack.py. Once approved,
merges DOWN into the layer below (a fast-forward); only the bottom layer
squash-merges into the real base. See ~/adocs/review-stack.md.
Rule: every changed file under tests-unit/ or tests/ (added, modified, or deleted)
Question: Is the code layer well covered, and did any edit weaken an existing check?
Source tip: eca2c74bff
Merge-base: 20d59d2a5f

* review-stack 7/8: ported-fixes (42 files, +1361/-180)

Review-and-land stack for synap5e/feat/assets-di-v2, generated by
review-stack.py conventions (hand-built continuation layer; see the PR body).
Once approved, merges DOWN into the layer below (a fast-forward); only the
bottom layer squash-merges into the real base. See ~/adocs/review-stack.md.
Rule: the 11 base-branch fix/docs commits 595cd6e4..94d7185b cherry-picked across the DI refactor (7efdd1d7 excluded, superseded by layer 8)
Question: was each base fix ported faithfully across the DI refactor?
Source tip: 6841881069284803b902b4a9e33bdcda13126771
Merge-base: 7fdfb40f4b

* review-stack 8/8: defensive-parity (4 files, +36/-3)

Review-and-land stack for synap5e/feat/assets-di-v2, generated by
review-stack.py conventions (hand-built continuation layer; see the PR body).
Once approved, merges DOWN into the layer below (a fast-forward); only the
bottom layer squash-merges into the real base. See ~/adocs/review-stack.md.
Rule: match-or-improve master's dependency defenses — NoAssets selection when DB deps unavailable (7efdd1d7's outcome via the DI seam), requirements warning before assets imports, blake3 in the guarded dependency set
Question: does each degradation path now match or improve master's behavior?
Source tip: ebc2cfeebc
Merge-base: 7fdfb40f4b

* fix(assets): only discard content rows this operation actually inserted

CR-9: Enumerated all six create_content call sites. Only scanner seeding and the three ingest registration paths track IDs for failure cleanup.

* fix(assets): reject hash-only uploads with FEATURE_DISABLED when hashing is off

CodeRabbit finding CR-2: reject hash-only multipart uploads before create_from_hash when hashing is disabled.

* fix(assets): seed persists the stat it verified

CR-7: persist the fresh seed-time restat instead of walk-time spec values.

* fix(assets): route database lock failures to the lock guidance

CR-16: route file-lock startup failures through the existing lock guidance and exit path.

* fix(assets): drop the inaccurate temp-cleanup claim from the shutdown warning

References CR-10.

* fix(assets): walk the output root after execution so undeclared outputs register promptly

Custom nodes that write files into the output directory without declaring
them in output_ui only became assets when the next full walk happened - a
frontend GET /object_info or a restart. Headless and API-only sessions never
trigger either, so those files never converged into the asset database.

The post-execution hook now requests a FULL scan of the output root instead of
an enrich-only pass. The seeder's pending-request queue was generalised from
enrich-specific to carrying a scan phase, so the request starts immediately
when the seeder is idle and coalesces (escalating to FULL on a phase mismatch)
when a scan is already running. queue_output_enrichment is renamed to
queue_output_scan across the protocol, the NoAssets no-op and the call site.

References FIX-6.

* chore(assets): remove seeder paths orphaned by the output-scan change

45c2f96e rerouted both former enrich call sites to start()/enqueue_scan(),
leaving two seeder methods that look live but are not. Review round F2
raised this along with four smaller items; the user's disposition was to
fix all six here.

- Delete start_enrich: zero callers repo-wide after 45c2f96e.
- Delete enqueue_enrich: no production callers; its ~18 call sites in
  tests/test_asset_seeder.py move to enqueue_scan(phase=ScanPhase.ENRICH)
  with their semantics unchanged. The deletion forces the half-done class
  renames (TestEnqueueEnrich* -> TestEnqueueScan*, consistent with the
  already-renamed TestPendingScanDrain) and restores the module docstring
  that was dropped rather than reworded.
- Document at manager.queue_output_scan that ScanPhase.FULL per debounce
  window is the deliberate, user-ratified trade, so it is not optimised
  back to ENRICH without revisiting the decision.
- Document that SeedAssetSpec.size_bytes/mtime_ns are walk-time
  diagnostics only - production persists the seed-time restat since CR-7.
- Export create_content_reporting_insert from the queries facade and fold
  scanner.py's direct-module import into the existing facade block.
- Harden test_queue_output_scan_does_not_duplicate_declared_output against
  a vacuous pass: it now asserts the seeder finished without errors and
  that an undeclared sibling written into the same directory WAS
  registered by the same scan, proving the walk actually ran.

No production behaviour changes beyond the two deletions.

References F2-cleanup.

* chore: comment cleanup

Comment-Gate: 18 quarantined

* fix(assets): preserve pause across the seeder's pending-scan drain

pause() runs before every prompt, while pending-scan enqueue and resume only run inside the debounced gc-interval gate. If the active scan finishes just after the next prompt's pause, its finally block resets the seeder to idle and the pending drain starts a replacement with the run gate open, so resume becomes a no-op.

Capture pausedness under the lock before resetting to idle, then start the drained scan already paused. Setting the state and gate before launching the thread avoids the start-then-reclear window and lets resume release the existing scan checkpoints.

Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>

* test(assets): pin job_id absence for scan-discovered assets

Owner ruling, recorded 2026-09-03 in the stack-9-hardening planning notepad: scan-discovered assets — including undeclared outputs found by the post-execution walk — carry job_id = None, always; only emission-time registration (output_ui declaration) attributes a job; attributing walk finds to the most recent prompt would be a temporal-correlation guess that is wrong exactly when prompts interleave; None is honest provenance. Do NOT add proximity-based attribution heuristics to the scanner. Ratified against Jacob Segal's cross-job-attribution concern (2026-09-08 review meeting) — a wrongly-attributed asset could mean one user's cloud job sees another user's asset.

Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)
Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>

* [review-stack 10/10] assets-tests (#16218)

* test(execution): run the battery with assets enabled and assert asset-system health at teardown

* test(execution): cover list-shaped outputs registering assets

Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>

---------

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>

* [review-stack 11/11] review-fixes (#16261)

* fix(assets): only exit on database file-lock timeout when assets are enabled

* test(assets): pin live_contents_under_prefixes path-filtering semantics

* perf(assets): push live-content prefix filtering into SQL

* test(assets): declare per-entry intent in the path-prefix corpus

* test(assets): normalize POSIX-literal path expectations for Windows

* test(assets): force observable stat changes and close-before-mutate on Windows-sensitive rewrites

* test(assets): force an observable mtime change in the hash-mode split test

---------

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
Co-authored-by: guill <jacob.e.segal@gmail.com>
This commit is contained in:
Simon Pinfold
2026-09-14 07:04:51 +12:00
committed by GitHub
parent 798fa9aa47
commit 19e1058f4c
132 changed files with 16452 additions and 13647 deletions

View File

@@ -150,7 +150,6 @@ def downgrade() -> None:
op.drop_index("ix_asset_cache_state_asset_id", table_name="asset_cache_state")
op.drop_index("ix_asset_cache_state_file_path", table_name="asset_cache_state")
op.drop_constraint("uq_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")
@@ -160,7 +159,6 @@ def downgrade() -> None:
op.drop_index("ix_tags_tag_type", table_name="tags")
op.drop_table("tags")
op.drop_constraint("uq_assets_info_asset_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")

View File

@@ -177,7 +177,10 @@ def downgrade() -> None:
NOTE: Data is not recoverable. The upgrade discards all rows from the old
tables and truncates assets. After downgrade the old schema will be empty.
A filesystem rescan will repopulate data once the older code is running.
A filesystem rescan can repopulate only data derived from paths. It will
not restore user_metadata, manually-applied tags, preview_id nominations,
name renames, or job_id, because build_asset_specs derives names and tags
from the path alone.
"""
# Drop new tables (order matters due to FK constraints)
op.drop_index("ix_asset_reference_meta_key_val_bool", table_name="asset_reference_meta")

View File

@@ -0,0 +1,125 @@
"""
Record/content split.
This migration intentionally discards the existing asset database. The
asset_reference_meta, asset_reference_tags, asset_references, and assets
tables are dropped, and DELETE FROM tags removes all existing tag rows. No
data migration is performed.
A filesystem rescan after this migration will not restore user_metadata,
manually-applied tags, preview_id nominations, name renames, or job_id.
build_asset_specs derives names and tags from the path alone.
Revision ID: 0007_record_content_split
Revises: 0006_add_loader_path
Create Date: 2026-08-26
"""
from alembic import op
import sqlalchemy as sa
revision = "0007_record_content_split"
down_revision = "0006_add_loader_path"
branch_labels = None
depends_on = None
def upgrade() -> None:
op.drop_table("asset_reference_meta")
op.drop_table("asset_reference_tags")
op.drop_table("asset_references")
op.drop_table("assets")
op.execute("DELETE FROM tags")
op.create_table(
"asset_contents",
sa.Column("id", sa.String(36), primary_key=True),
sa.Column("hash", sa.String(256)),
sa.Column("size_bytes", sa.BigInteger(), nullable=False),
sa.Column("path", sa.Text(), nullable=False),
sa.Column("mtime_ns", sa.BigInteger()),
sa.Column("is_missing", sa.Boolean(), nullable=False, server_default="0"),
sa.Column("created_at", sa.DateTime(), nullable=False),
sa.CheckConstraint("size_bytes >= 0", name="ck_asset_contents_size_nonneg"),
sa.CheckConstraint("mtime_ns >= 0", name="ck_asset_contents_mtime_nonneg"),
)
op.create_index("ix_asset_contents_hash", "asset_contents", ["hash"])
op.create_index(
"uq_asset_contents_path_live", "asset_contents", ["path"], unique=True,
sqlite_where=sa.text("is_missing = 0"),
)
op.create_table(
"assets",
sa.Column("id", sa.String(36), primary_key=True),
sa.Column("content_id", sa.String(36), sa.ForeignKey("asset_contents.id", ondelete="RESTRICT"), nullable=False),
sa.Column("name", sa.String(512), nullable=False),
sa.Column("mime_type", sa.String(255)),
sa.Column("system_metadata", sa.JSON()),
sa.Column("job_id", sa.String(36)),
sa.Column("user_metadata", sa.JSON()),
sa.Column("loader_path", sa.Text()),
sa.Column("preview_id", sa.String(36), sa.ForeignKey("assets.id", ondelete="SET NULL")),
sa.Column("created_at", sa.DateTime(), nullable=False),
sa.Column("updated_at", sa.DateTime(), nullable=False),
sa.Column("last_access_time", sa.DateTime()),
)
op.create_index("ix_assets_content_id", "assets", ["content_id"])
op.create_index("ix_assets_name", "assets", ["name"])
op.create_index("ix_assets_created_at", "assets", ["created_at"])
op.create_index("ix_assets_preview_id", "assets", ["preview_id"])
op.create_table(
"asset_meta",
sa.Column("asset_id", sa.String(36), sa.ForeignKey("assets.id", ondelete="CASCADE"), primary_key=True),
sa.Column("key", sa.String(256), primary_key=True),
sa.Column("ordinal", sa.Integer(), primary_key=True),
sa.Column("val_str", sa.String(2048)), sa.Column("val_num", sa.Numeric(38, 10)),
sa.Column("val_bool", sa.Boolean()), sa.Column("val_json", sa.JSON()),
sa.CheckConstraint("val_str IS NOT NULL OR val_num IS NOT NULL OR val_bool IS NOT NULL OR val_json IS NOT NULL", name="ck_asset_meta_has_value"),
)
op.create_index("ix_asset_meta_key", "asset_meta", ["key"])
op.create_index("ix_asset_meta_key_val_str", "asset_meta", ["key", "val_str"])
op.create_index("ix_asset_meta_key_val_num", "asset_meta", ["key", "val_num"])
op.create_index("ix_asset_meta_key_val_bool", "asset_meta", ["key", "val_bool"])
op.create_table("asset_tags", sa.Column("asset_id", sa.String(36), sa.ForeignKey("assets.id", ondelete="CASCADE"), primary_key=True), sa.Column("tag_name", sa.String(512), sa.ForeignKey("tags.name", ondelete="RESTRICT"), primary_key=True), sa.Column("origin", sa.String(32), nullable=False), sa.Column("added_at", sa.DateTime(), nullable=False))
op.create_index("ix_asset_tags_tag_name", "asset_tags", ["tag_name"])
op.create_index("ix_asset_tags_asset_id", "asset_tags", ["asset_id"])
op.create_table("asset_system_state", sa.Column("key", sa.String(256), primary_key=True), sa.Column("value", sa.Text(), nullable=False))
def downgrade() -> None:
op.drop_table("asset_system_state")
op.drop_table("asset_tags")
op.drop_table("asset_meta")
op.drop_table("assets")
op.drop_table("asset_contents")
op.create_table(
"assets",
sa.Column("id", sa.String(length=36), primary_key=True),
sa.Column("hash", sa.String(length=256), nullable=True),
sa.Column("size_bytes", sa.BigInteger(), nullable=False, server_default="0"),
sa.Column("mime_type", sa.String(length=255), nullable=True),
sa.Column("created_at", sa.DateTime(timezone=False), nullable=False),
sa.CheckConstraint("size_bytes >= 0", name="ck_assets_size_nonneg"),
)
op.create_index("uq_assets_hash", "assets", ["hash"], unique=True)
op.create_index("ix_assets_mime_type", "assets", ["mime_type"])
op.create_table("asset_references", sa.Column("id", sa.String(36), primary_key=True), sa.Column("asset_id", sa.String(36), sa.ForeignKey("assets.id", ondelete="CASCADE"), nullable=False), sa.Column("file_path", sa.Text()), sa.Column("loader_path", sa.Text()), sa.Column("mtime_ns", sa.BigInteger()), sa.Column("needs_verify", sa.Boolean(), nullable=False), sa.Column("is_missing", sa.Boolean(), nullable=False), sa.Column("enrichment_level", sa.Integer(), nullable=False), sa.Column("owner_id", sa.String(128), nullable=False), sa.Column("name", sa.String(512), nullable=False), sa.Column("preview_id", sa.String(36), sa.ForeignKey("asset_references.id", ondelete="SET NULL")), sa.Column("user_metadata", sa.JSON()), sa.Column("system_metadata", sa.JSON()), sa.Column("job_id", sa.String(36)), sa.Column("created_at", sa.DateTime(), nullable=False), sa.Column("updated_at", sa.DateTime(), nullable=False), sa.Column("last_access_time", sa.DateTime(), nullable=False), sa.Column("deleted_at", sa.DateTime()))
op.create_index("uq_asset_references_file_path", "asset_references", ["file_path"], unique=True)
op.create_index("ix_asset_references_asset_id", "asset_references", ["asset_id"])
op.create_index("ix_asset_references_owner_id", "asset_references", ["owner_id"])
op.create_index("ix_asset_references_name", "asset_references", ["name"])
op.create_index("ix_asset_references_is_missing", "asset_references", ["is_missing"])
op.create_index("ix_asset_references_enrichment_level", "asset_references", ["enrichment_level"])
op.create_index("ix_asset_references_created_at", "asset_references", ["created_at"])
op.create_index("ix_asset_references_last_access_time", "asset_references", ["last_access_time"])
op.create_index("ix_asset_references_owner_name", "asset_references", ["owner_id", "name"])
op.create_index("ix_asset_references_deleted_at", "asset_references", ["deleted_at"])
op.create_index("ix_asset_references_preview_id", "asset_references", ["preview_id"])
op.create_table("asset_reference_meta", sa.Column("asset_reference_id", sa.String(36), sa.ForeignKey("asset_references.id", ondelete="CASCADE"), primary_key=True), sa.Column("key", sa.String(256), primary_key=True), sa.Column("ordinal", sa.Integer(), primary_key=True), sa.Column("val_str", sa.String(2048)), sa.Column("val_num", sa.Numeric(38, 10)), sa.Column("val_bool", sa.Boolean()), sa.Column("val_json", sa.JSON()), sa.CheckConstraint("val_str IS NOT NULL OR val_num IS NOT NULL OR val_bool IS NOT NULL OR val_json IS NOT NULL", name="ck_asset_reference_meta_has_value"))
op.create_index("ix_asset_reference_meta_key", "asset_reference_meta", ["key"])
op.create_index("ix_asset_reference_meta_key_val_str", "asset_reference_meta", ["key", "val_str"])
op.create_index("ix_asset_reference_meta_key_val_num", "asset_reference_meta", ["key", "val_num"])
op.create_index("ix_asset_reference_meta_key_val_bool", "asset_reference_meta", ["key", "val_bool"])
op.create_table("asset_reference_tags", sa.Column("asset_reference_id", sa.String(36), sa.ForeignKey("asset_references.id", ondelete="CASCADE"), primary_key=True), sa.Column("tag_name", sa.String(512), sa.ForeignKey("tags.name", ondelete="RESTRICT"), primary_key=True), sa.Column("origin", sa.String(32), nullable=False), sa.Column("added_at", sa.DateTime(), nullable=False))
op.create_index("ix_asset_reference_tags_tag_name", "asset_reference_tags", ["tag_name"])
op.create_index("ix_asset_reference_tags_asset_reference_id", "asset_reference_tags", ["asset_reference_id"])

View File

@@ -1,3 +1,11 @@
"""Serves the asset HTTP API: listing, detail, upload, tagging, hash-addressed
creation and scan control. A feature gate wraps every handler and answers 503
when the asset system is switched off or its database dependencies are absent,
and failures leave as a JSON envelope carrying a stable machine-readable code
rather than aiohttp's default text. Handlers own request parsing and error
shaping only; the work itself belongs to the services layer.
"""
import asyncio
import functools
import json
@@ -6,6 +14,7 @@ import mimetypes
import os
import urllib.parse
import uuid
from datetime import timezone
from typing import Any
from aiohttp import web
@@ -13,6 +22,7 @@ from pydantic import ValidationError
import folder_paths
from app import user_manager
from app.assets import mode
from app.assets.api import schemas_in, schemas_out
from app.assets.services import schemas
from app.assets.api.schemas_in import (
@@ -24,30 +34,54 @@ from app.assets.api.upload import (
delete_temp_file_if_exists,
parse_multipart_upload,
)
from app.assets.database.models import Asset
from app.assets.database.queries.records import (
RecordCursorBoundary,
RecordPageSpec,
RecordSortField,
RecordSortOrder,
get_preview_file_paths_by_ids,
list_records_page,
)
from app.assets.seeder import ScanInProgressError, asset_seeder
from app.assets.services import (
DependencyMissingError,
HashMismatchError,
UploadUnstableError,
apply_tags,
asset_exists,
create_from_hash,
delete_asset_reference,
get_asset_detail,
get_preview_file_paths,
list_assets_page,
list_tags,
remove_tags,
resolve_asset_for_download,
update_asset_metadata,
upload_from_temp_path,
)
from app.assets.services.cursor import InvalidCursorError
from app.assets.services.path_utils import compute_asset_response_paths
from app.assets.services.cursor import (
InvalidCursorError,
decode_cursor,
decode_cursor_int,
decode_cursor_time,
encode_cursor,
encode_cursor_from_time,
)
from app.assets.services.tagging import list_tag_histogram
from app.database.db import create_session
ROUTES = web.RouteTableDef()
USER_MANAGER: user_manager.UserManager | None = None
_ASSETS_ENABLED = False
SYSTEM_TAGS = frozenset({"missing"})
_CURSOR_SORT_FIELDS: tuple[RecordSortField, ...] = (
"created_at",
"updated_at",
"name",
"size",
)
def _require_assets_feature_enabled(handler):
@@ -99,12 +133,6 @@ def register_assets_routes(
app.add_routes(ROUTES)
def disable_assets_routes() -> None:
"""Disable asset routes at runtime (e.g. after DB init failure)."""
global _ASSETS_ENABLED
_ASSETS_ENABLED = False
def _build_error_response(
status: int, code: str, message: str, details: dict | None = None
) -> web.Response:
@@ -119,6 +147,20 @@ def _build_validation_error_response(code: str, ve: ValidationError) -> web.Resp
return _build_error_response(400, code, "Validation failed.", {"errors": errors})
class SystemTagForbiddenError(Exception):
def __init__(self, tag: str):
super().__init__(
f"Tag '{tag}' is system-managed and cannot be modified via the API"
)
self.tag = tag
def _reject_system_tags(tags: list[str]) -> None:
for tag in tags:
if tag in SYSTEM_TAGS:
raise SystemTagForbiddenError(tag)
class InvalidTagFilterError(Exception):
"""Invalid combination of tag-filter query parameters."""
@@ -200,13 +242,22 @@ def _resolve_tag_filters(
return all_list, tags_any, none_list
def _validate_sort_field(requested: str | None) -> str:
def _validate_sort_field(requested: str | None) -> RecordSortField:
if not requested:
return "created_at"
v = requested.lower()
if v in {"name", "created_at", "updated_at", "size", "last_access_time"}:
return v
return "created_at"
match requested.lower():
case "name":
return "name"
case "created_at":
return "created_at"
case "updated_at":
return "updated_at"
case "size":
return "size"
case "last_access_time":
return "last_access_time"
case _:
return "created_at"
# What a client can render from the bytes themselves; anything else needs a nominated preview.
@@ -258,6 +309,8 @@ def _build_asset_response(
if result.ref.preview_id:
# A nominated preview is one whatever it holds, so no media check here.
preview_url = _build_view_url(preview_paths.get(result.ref.preview_id))
elif result.asset is not None and result.asset.is_missing:
preview_url = None
elif _has_previewable_content(result.asset, result.ref.file_path):
preview_url = _build_view_url(result.ref.file_path)
else:
@@ -268,7 +321,8 @@ def _build_asset_response(
# In-root loader path (model category dropped): what model loaders consume.
loader_path = result.ref.loader_path
else:
display_name, loader_path = None, None
display_name = None
loader_path = None
asset_content_hash = result.asset.hash if result.asset else None
return schemas_out.Asset(
id=result.ref.id,
@@ -276,7 +330,6 @@ def _build_asset_response(
hash=asset_content_hash,
loader_path=loader_path,
display_name=display_name,
asset_hash=asset_content_hash,
size=int(result.asset.size_bytes) if result.asset else None,
mime_type=result.asset.mime_type if result.asset else None,
tags=result.tags,
@@ -292,6 +345,111 @@ def _build_asset_response(
)
def _build_record_response(
record: Asset,
tags: list[str],
preview_paths: dict[str, str],
) -> schemas_out.Asset:
content = record.content
paths = compute_asset_response_paths(content.path)
display_name = paths[1] if paths else None
if record.preview_id:
preview_url = _build_view_url(preview_paths.get(record.preview_id))
elif content.is_missing:
preview_url = None
else:
mime_type = record.mime_type or mimetypes.guess_type(content.path)[0] or ""
if mime_type.split(";", 1)[0].strip().lower().startswith(
PREVIEWABLE_MIME_PREFIXES
):
preview_url = _build_view_url(content.path)
else:
preview_url = None
return schemas_out.Asset(
id=record.id,
name=record.name,
hash=content.hash,
loader_path=record.loader_path,
display_name=display_name,
size=content.size_bytes,
mime_type=record.mime_type,
tags=tags,
preview_url=preview_url,
preview_id=record.preview_id,
user_metadata=record.user_metadata or {},
metadata=record.system_metadata,
job_id=record.job_id,
prompt_id=record.job_id,
created_at=record.created_at,
updated_at=record.updated_at,
last_access_time=record.last_access_time,
)
def _decode_record_cursor(
after: str | None,
sort: RecordSortField,
order: RecordSortOrder,
) -> RecordCursorBoundary | None:
if after is None:
return None
if sort not in _CURSOR_SORT_FIELDS:
raise InvalidCursorError(
f"cursor pagination is not supported for sort={sort!r}"
)
payload = decode_cursor(
after,
_CURSOR_SORT_FIELDS,
expected_order=order,
)
if payload.sort_field != sort:
raise InvalidCursorError(
f"cursor sort field {payload.sort_field!r} does not match request sort {sort!r}"
)
match payload.sort_field:
case "created_at" | "updated_at":
value = decode_cursor_time(payload).replace(tzinfo=None)
case "size":
value = decode_cursor_int(payload)
case "name":
value = payload.value
case unsupported:
raise InvalidCursorError(f"unsupported sort field {unsupported!r}")
return RecordCursorBoundary(value=value, id=payload.id)
def _encode_record_cursor(
record: Asset,
sort: RecordSortField,
order: RecordSortOrder,
) -> str:
match sort:
case "name":
return encode_cursor("name", record.name, record.id, order=order)
case "size":
return encode_cursor(
"size",
str(record.content.size_bytes),
record.id,
order=order,
)
case "created_at":
timestamp = record.created_at
case "updated_at":
timestamp = record.updated_at
case "last_access_time":
raise InvalidCursorError(
"cursor pagination is not supported for sort='last_access_time'"
)
return encode_cursor_from_time(
sort,
timestamp.replace(tzinfo=timezone.utc),
record.id,
order=order,
)
@ROUTES.head("/api/assets/hash/{hash}")
@_require_assets_feature_enabled
async def head_asset_by_hash(request: web.Request) -> web.Response:
@@ -312,6 +470,11 @@ async def list_assets_route(request: web.Request) -> web.Response:
"""
GET request to list assets.
"""
if "metadata_filter" in request.query:
return _build_error_response(
400, "UNSUPPORTED_PARAM", "metadata_filter is no longer supported"
)
query_dict = get_query_dict(request)
try:
q = schemas_in.ListAssetsQuery.model_validate(query_dict)
@@ -324,42 +487,57 @@ async def list_assets_route(request: web.Request) -> web.Response:
return _build_error_response(400, "INVALID_TAG_FILTER", str(e), e.details)
sort = _validate_sort_field(q.sort)
order_candidate = (q.order or "desc").lower()
order = order_candidate if order_candidate in {"asc", "desc"} else "desc"
order = q.order
try:
result = list_assets_page(
owner_id=USER_MANAGER.get_request_user_id(request),
include_tags=tags_all,
exclude_tags=tags_none,
any_tags=tags_any,
name_contains=q.name_contains,
metadata_filter=q.metadata_filter,
limit=q.limit,
offset=q.offset,
sort=sort,
order=order,
after=q.after,
)
cursor_boundary = _decode_record_cursor(q.after, sort, order)
cursor_supported = sort in _CURSOR_SORT_FIELDS
fetch_limit = q.limit + 1 if cursor_supported else q.limit
with create_session() as session:
records, tag_map, total = list_records_page(
session,
RecordPageSpec(
all_tags=tuple(tags_all),
any_tags=tuple(tags_any),
none_tags=tuple(tags_none),
name_contains=q.name_contains,
limit=fetch_limit,
offset=q.offset,
sort=sort,
order=order,
after=cursor_boundary,
),
)
next_cursor = None
if cursor_supported and len(records) > q.limit:
records = records[: q.limit]
next_cursor = _encode_record_cursor(records[-1], sort, order)
preview_ids = sorted(
{record.preview_id for record in records if record.preview_id}
)
preview_paths = get_preview_file_paths_by_ids(session, preview_ids)
summaries = [
_build_record_response(
record,
tag_map.get(record.id, []),
preview_paths,
)
for record in records
]
except InvalidCursorError as e:
return _build_error_response(400, "INVALID_CURSOR", str(e))
preview_paths = _resolve_preview_paths(result.items)
summaries = [_build_asset_response(item, preview_paths) for item in result.items]
# has_more semantics differ by mode:
# - cursor mode: a non-empty next_cursor means there are more results.
# - offset mode: derived from total - (offset + page size).
if q.after is not None:
has_more = result.next_cursor is not None
has_more = next_cursor is not None
else:
has_more = (q.offset + len(summaries)) < result.total
has_more = q.offset + len(summaries) < total
payload = schemas_out.AssetsList(
assets=summaries,
total=result.total,
total=total,
has_more=has_more,
next_cursor=result.next_cursor,
next_cursor=next_cursor,
)
return web.json_response(payload.model_dump(mode="json", exclude_none=True))
@@ -374,7 +552,6 @@ async def get_asset_route(request: web.Request) -> web.Response:
try:
result = get_asset_detail(
reference_id=reference_id,
owner_id=USER_MANAGER.get_request_user_id(request),
)
if not result:
return _build_error_response(
@@ -391,7 +568,7 @@ async def get_asset_route(request: web.Request) -> web.Response:
)
except Exception:
logging.exception(
"get_asset failed for reference_id=%s, owner_id=%s",
"get_asset failed for reference_id=%s, tenant_id=%s",
reference_id,
USER_MANAGER.get_request_user_id(request),
)
@@ -409,7 +586,6 @@ async def download_asset_content(request: web.Request) -> web.Response:
try:
result = resolve_asset_for_download(
reference_id=str(uuid.UUID(request.match_info["id"])),
owner_id=USER_MANAGER.get_request_user_id(request),
)
abs_path = result.abs_path
content_type = result.content_type
@@ -501,15 +677,22 @@ async def create_asset_from_hash_route(request: web.Request) -> web.Response:
if name is None:
name = body.hash.split(":", 1)[1] if ":" in body.hash else body.hash
result = create_from_hash(
hash_str=body.hash,
name=name,
tags=body.tags,
user_metadata=body.user_metadata,
owner_id=USER_MANAGER.get_request_user_id(request),
mime_type=body.mime_type,
preview_id=body.preview_id,
)
if not mode.hashing_enabled():
return _build_error_response(
400, "FEATURE_DISABLED", "Asset hashing is disabled."
)
try:
result = create_from_hash(
hash_str=body.hash,
name=name,
tags=body.tags,
user_metadata=body.user_metadata,
mime_type=body.mime_type,
preview_id=body.preview_id,
)
except ValueError as e:
return _build_error_response(400, "INVALID_BODY", str(e))
if result is None:
return _build_error_response(
404, "ASSET_NOT_FOUND", f"Asset content {body.hash} does not exist"
@@ -532,7 +715,7 @@ async def upload_asset(request: web.Request) -> web.Response:
except UploadError as e:
return _build_error_response(e.status, e.code, e.message)
owner_id = USER_MANAGER.get_request_user_id(request)
tenant_id = USER_MANAGER.get_request_user_id(request)
try:
spec = schemas_in.UploadAssetSpec.model_validate(
@@ -552,43 +735,40 @@ async def upload_asset(request: web.Request) -> web.Response:
)
try:
# Fast path: hash exists, create AssetReference without writing anything
if spec.hash and parsed.provided_hash_exists is True:
if not parsed.file_present and spec.hash:
if not mode.hashing_enabled():
return _build_error_response(
400, "FEATURE_DISABLED", "Asset hashing is disabled."
)
result = create_from_hash(
hash_str=spec.hash,
name=spec.name or (spec.hash.split(":", 1)[1]),
tags=spec.tags,
user_metadata=spec.user_metadata or {},
owner_id=owner_id,
mime_type=spec.mime_type,
preview_id=spec.preview_id,
)
if result is None:
delete_temp_file_if_exists(parsed.tmp_path)
return _build_error_response(
404, "ASSET_NOT_FOUND", f"Asset content {spec.hash} does not exist"
)
delete_temp_file_if_exists(parsed.tmp_path)
else:
# Otherwise, we must have a temp file path to ingest
if not parsed.tmp_path or not os.path.exists(parsed.tmp_path):
return _build_error_response(
400,
"MISSING_INPUT",
"Provided hash not found and no file uploaded.",
)
elif parsed.tmp_path and os.path.exists(parsed.tmp_path):
result = upload_from_temp_path(
temp_path=parsed.tmp_path,
name=spec.name,
tags=spec.tags,
user_metadata=spec.user_metadata or {},
client_filename=parsed.file_client_name,
owner_id=owner_id,
expected_hash=spec.hash,
mime_type=spec.mime_type,
preview_id=spec.preview_id,
)
else:
return _build_error_response(
400,
"MISSING_INPUT",
"Provided hash not found and no file uploaded.",
)
except AssetValidationError as e:
delete_temp_file_if_exists(parsed.tmp_path)
return _build_error_response(400, e.code, str(e))
@@ -598,12 +778,15 @@ async def upload_asset(request: web.Request) -> web.Response:
except HashMismatchError as e:
delete_temp_file_if_exists(parsed.tmp_path)
return _build_error_response(400, "HASH_MISMATCH", str(e))
except UploadUnstableError as e:
delete_temp_file_if_exists(parsed.tmp_path)
return _build_error_response(500, "UPLOAD_UNSTABLE", str(e))
except DependencyMissingError as e:
delete_temp_file_if_exists(parsed.tmp_path)
return _build_error_response(503, "DEPENDENCY_MISSING", e.message)
except Exception:
delete_temp_file_if_exists(parsed.tmp_path)
logging.exception("upload_asset failed for owner_id=%s", owner_id)
logging.exception("upload_asset failed for tenant_id=%s", tenant_id)
return _build_error_response(500, "INTERNAL", "Unexpected server error.")
asset = _build_asset_response(result, _resolve_preview_paths([result]))
@@ -633,7 +816,6 @@ async def update_asset_route(request: web.Request) -> web.Response:
reference_id=reference_id,
name=body.name,
user_metadata=body.user_metadata,
owner_id=USER_MANAGER.get_request_user_id(request),
preview_id=body.preview_id,
)
payload = _build_asset_response(result, _resolve_preview_paths([result]))
@@ -645,7 +827,7 @@ async def update_asset_route(request: web.Request) -> web.Response:
)
except Exception:
logging.exception(
"update_asset failed for reference_id=%s, owner_id=%s",
"update_asset failed for reference_id=%s, tenant_id=%s",
reference_id,
USER_MANAGER.get_request_user_id(request),
)
@@ -659,16 +841,12 @@ async def delete_asset_route(request: web.Request) -> web.Response:
reference_id = str(uuid.UUID(request.match_info["id"]))
try:
# Deleting an asset is a soft delete of the reference; the underlying
# content is preserved (it may be shared with other references).
deleted = delete_asset_reference(
reference_id=reference_id,
owner_id=USER_MANAGER.get_request_user_id(request),
delete_content_if_orphan=False,
)
except Exception:
logging.exception(
"delete_asset_reference failed for reference_id=%s, owner_id=%s",
"delete_asset_reference failed for reference_id=%s, tenant_id=%s",
reference_id,
USER_MANAGER.get_request_user_id(request),
)
@@ -705,7 +883,6 @@ async def get_tags(request: web.Request) -> web.Response:
offset=query.offset,
order=query.order,
include_zero=query.include_zero,
owner_id=USER_MANAGER.get_request_user_id(request),
)
tags = [
@@ -738,17 +915,21 @@ async def add_asset_tags(request: web.Request) -> web.Response:
)
try:
_reject_system_tags(data.tags)
result = apply_tags(
reference_id=reference_id,
tags=data.tags,
origin="manual",
owner_id=USER_MANAGER.get_request_user_id(request),
)
payload = schemas_out.TagsAdd(
added=result.added,
already_present=result.already_present,
total_tags=result.total_tags,
)
except SystemTagForbiddenError as se:
return _build_error_response(
400, "SYSTEM_TAG_FORBIDDEN", str(se), {"tag": se.tag}
)
except PermissionError as pe:
return _build_error_response(403, "FORBIDDEN", str(pe), {"id": reference_id})
except ValueError as ve:
@@ -757,7 +938,7 @@ async def add_asset_tags(request: web.Request) -> web.Response:
)
except Exception:
logging.exception(
"add_tags_to_asset failed for reference_id=%s, owner_id=%s",
"add_tags_to_asset failed for reference_id=%s, tenant_id=%s",
reference_id,
USER_MANAGER.get_request_user_id(request),
)
@@ -786,15 +967,20 @@ async def delete_asset_tags(request: web.Request) -> web.Response:
)
try:
_reject_system_tags(data.tags)
result = remove_tags(
reference_id=reference_id,
tags=data.tags,
owner_id=USER_MANAGER.get_request_user_id(request),
)
payload = schemas_out.TagsRemove(
removed=result.removed,
not_present=result.not_present,
total_tags=result.total_tags,
protected=result.protected,
)
except SystemTagForbiddenError as se:
return _build_error_response(
400, "SYSTEM_TAG_FORBIDDEN", str(se), {"tag": se.tag}
)
except PermissionError as pe:
return _build_error_response(403, "FORBIDDEN", str(pe), {"id": reference_id})
@@ -804,7 +990,7 @@ async def delete_asset_tags(request: web.Request) -> web.Response:
)
except Exception:
logging.exception(
"remove_tags_from_asset failed for reference_id=%s, owner_id=%s",
"remove_tags_from_asset failed for reference_id=%s, tenant_id=%s",
reference_id,
USER_MANAGER.get_request_user_id(request),
)
@@ -817,6 +1003,11 @@ async def delete_asset_tags(request: web.Request) -> web.Response:
@_require_assets_feature_enabled
async def get_tags_refine(request: web.Request) -> web.Response:
"""GET request to get tag histogram for filtered assets."""
if "metadata_filter" in request.query:
return _build_error_response(
400, "UNSUPPORTED_PARAM", "metadata_filter is no longer supported"
)
query_dict = get_query_dict(request)
try:
q = schemas_in.TagsRefineQuery.model_validate(query_dict)
@@ -829,12 +1020,10 @@ async def get_tags_refine(request: web.Request) -> web.Response:
return _build_error_response(400, "INVALID_TAG_FILTER", str(e), e.details)
tag_counts = list_tag_histogram(
owner_id=USER_MANAGER.get_request_user_id(request),
include_tags=tags_all,
exclude_tags=tags_none,
any_tags=tags_any,
name_contains=q.name_contains,
metadata_filter=q.metadata_filter,
limit=q.limit,
)
payload = schemas_out.TagHistogram(tag_counts=tag_counts)
@@ -867,7 +1056,9 @@ async def seed_assets(request: web.Request) -> web.Response:
wait_param = request.query.get("wait", "").lower()
should_wait = wait_param in ("true", "1", "yes")
started = asset_seeder.start(roots=valid_roots)
started = asset_seeder.start(
roots=valid_roots, compute_hashes=mode.hashing_enabled()
)
if not started:
return web.json_response({"status": "already_running"}, status=409)

View File

@@ -1,3 +1,10 @@
"""Parses and validates everything the asset API accepts from a client, so
handlers receive typed values instead of raw JSON. Query strings, JSON bodies
and multipart upload specs each get a model that rejects malformed input at the
boundary, normalizes tags and hashes, and raises errors already carrying the
HTTP status and code the handler should return.
"""
import json
from dataclasses import dataclass
from typing import Any, Literal
@@ -58,9 +65,6 @@ class ListAssetsQuery(BaseModel):
tags_none: list[str] = Field(default_factory=list)
name_contains: str | None = None
# Accept either a JSON string (query param) or a dict
metadata_filter: dict[str, Any] | None = None
limit: conint(ge=1, le=500) = 20
offset: conint(ge=0) = 0
# Opaque keyset cursor. When supplied, `offset` is ignored. Cursor pagination
@@ -93,22 +97,6 @@ class ListAssetsQuery(BaseModel):
return out
return v
@field_validator("metadata_filter", mode="before")
@classmethod
def _parse_metadata_json(cls, v):
if v is None or isinstance(v, dict):
return v
if isinstance(v, str) and v.strip():
try:
parsed = json.loads(v)
except Exception as e:
raise ValueError(f"metadata_filter must be JSON: {e}") from e
if not isinstance(parsed, dict):
raise ValueError("metadata_filter must be a JSON object")
return parsed
return None
class UpdateAssetBody(BaseModel):
name: str | None = None
user_metadata: dict[str, Any] | None = None
@@ -168,7 +156,6 @@ class TagsRefineQuery(BaseModel):
tags_any: list[str] = Field(default_factory=list)
tags_none: list[str] = Field(default_factory=list)
name_contains: str | None = None
metadata_filter: dict[str, Any] | None = None
limit: conint(ge=1, le=1000) = 100
@field_validator(
@@ -189,22 +176,6 @@ class TagsRefineQuery(BaseModel):
return out
return v
@field_validator("metadata_filter", mode="before")
@classmethod
def _parse_metadata_json(cls, v):
if v is None or isinstance(v, dict):
return v
if isinstance(v, str) and v.strip():
try:
parsed = json.loads(v)
except Exception as e:
raise ValueError(f"metadata_filter must be JSON: {e}") from e
if not isinstance(parsed, dict):
raise ValueError("metadata_filter must be a JSON object")
return parsed
return None
class TagsListQuery(BaseModel):
model_config = ConfigDict(extra="ignore", str_strip_whitespace=True)

View File

@@ -1,3 +1,9 @@
"""Defines the response bodies the asset API returns, fixing the field names,
types and serialization its HTTP contract promises. Routes dump these models
with null fields omitted, so whether a field is absent and whether it is empty
stay distinguishable to a client.
"""
from datetime import datetime
from typing import Any
@@ -5,14 +11,10 @@ from pydantic import BaseModel, ConfigDict, Field, field_serializer
class Asset(BaseModel):
"""API view of an asset. Maps to DB ``AssetReference`` joined with its ``Asset`` blob;
``id`` here is the AssetReference id, not the content-addressed Asset id."""
id: str
name: str = Field(
...,
deprecated=True,
description="Reference label, often caller-provided or derived from the filename. Deprecated for storage path/display semantics; use `loader_path` and `display_name` when present.",
description="Record label, usually derived from the source filename.",
)
hash: str | None = None
loader_path: str | None = Field(
@@ -23,14 +25,13 @@ class Asset(BaseModel):
default=None,
description="Human-facing label for the asset. Not unique.",
)
asset_hash: str | None = None
is_immutable: bool = False
size: int | None = None
mime_type: str | None = None
tags: list[str] = Field(default_factory=list)
preview_url: str | None = None
preview_id: str | None = None # references an asset_reference id, not an asset id
user_metadata: dict[str, Any] = Field(default_factory=dict)
is_immutable: bool = False
metadata: dict[str, Any] | None = None
job_id: str | None = None
prompt_id: str | None = None # deprecated: use job_id
@@ -80,6 +81,7 @@ class TagsRemove(BaseModel):
removed: list[str] = Field(default_factory=list)
not_present: list[str] = Field(default_factory=list)
total_tags: list[str] = Field(default_factory=list)
protected: list[str] = Field(default_factory=list)
class TagHistogram(BaseModel):

View File

@@ -1,3 +1,11 @@
"""Reads a multipart upload off the wire and lands its bytes in a temporary file
the ingest service can hash and move. The body is consumed in chunks so a large
model never has to fit in memory, and the temporary file is removed on every
failure path so an abandoned upload leaves nothing behind. Field values are
validated as they arrive, letting a bad request fail before its bytes are
written.
"""
import logging
import os
import uuid
@@ -91,20 +99,6 @@ async def parse_multipart_upload(
file_present = True
file_client_name = (field.filename or "").strip()
if provided_hash and provided_hash_exists is True:
# Hash exists - drain file but don't write to disk
try:
while True:
chunk = await field.read_chunk(8 * 1024 * 1024)
if not chunk:
break
file_written += len(chunk)
except Exception:
raise UploadError(
500, "UPLOAD_IO_ERROR", "Failed to receive uploaded file."
)
continue
uploads_root = os.path.join(folder_paths.get_temp_directory(), "uploads")
unique_dir = os.path.join(uploads_root, uuid.uuid4().hex)
os.makedirs(unique_dir, exist_ok=True)

View File

@@ -1,7 +1,16 @@
"""Declares the asset schema: content rows describing bytes on disk, records
describing what a user sees, and the tag and metadata tables hanging off them.
The split is the point — many records can name one content row, and retiring
content by marking it missing rather than deleting it is what keeps a path's
history intact. Constraints declared here, not application code, are what make
a negative size or a second live row at one path impossible.
"""
from __future__ import annotations
import uuid
from datetime import datetime
from decimal import Decimal
from typing import Any
from sqlalchemy import (
@@ -16,197 +25,134 @@ from sqlalchemy import (
Numeric,
String,
Text,
text,
)
from sqlalchemy.orm import Mapped, foreign, mapped_column, relationship
from sqlalchemy.orm import Mapped, mapped_column, relationship
from app.assets.helpers import get_utc_now
from app.database.models import Base
class AssetContent(Base):
__tablename__ = "asset_contents"
id: Mapped[str] = mapped_column(
String(36), primary_key=True, default=lambda: str(uuid.uuid4())
)
hash: Mapped[str | None] = mapped_column(String(256), index=True)
size_bytes: Mapped[int] = mapped_column(
BigInteger,
CheckConstraint("size_bytes >= 0", name="ck_asset_contents_size_nonneg"),
nullable=False,
default=0,
)
path: Mapped[str] = mapped_column(Text, nullable=False)
mtime_ns: Mapped[int | None] = mapped_column(
BigInteger,
CheckConstraint("mtime_ns >= 0", name="ck_asset_contents_mtime_nonneg"),
)
is_missing: Mapped[bool] = mapped_column(
Boolean, nullable=False, server_default="0"
)
created_at: Mapped[datetime] = mapped_column(
DateTime(timezone=False), nullable=False, default=get_utc_now
)
records: Mapped[list[Asset]] = relationship(back_populates="content")
__table_args__ = (
Index(
"uq_asset_contents_path_live",
"path",
unique=True,
sqlite_where=text("is_missing = 0"),
),
)
class Asset(Base):
__tablename__ = "assets"
id: Mapped[str] = mapped_column(
String(36), primary_key=True, default=lambda: str(uuid.uuid4())
)
hash: Mapped[str | None] = mapped_column(String(256), nullable=True)
size_bytes: Mapped[int] = mapped_column(BigInteger, nullable=False, default=0)
mime_type: Mapped[str | None] = mapped_column(String(255))
created_at: Mapped[datetime] = mapped_column(
DateTime(timezone=False), nullable=False, default=get_utc_now
content_id: Mapped[str] = mapped_column(
String(36),
ForeignKey("asset_contents.id", ondelete="RESTRICT"),
nullable=False,
)
references: Mapped[list[AssetReference]] = relationship(
"AssetReference",
back_populates="asset",
primaryjoin=lambda: Asset.id == foreign(AssetReference.asset_id),
foreign_keys=lambda: [AssetReference.asset_id],
cascade="all,delete-orphan",
passive_deletes=True,
)
# preview_id on AssetReference is a self-referential FK to asset_references.id
__table_args__ = (
Index("uq_assets_hash", "hash", unique=True),
Index("ix_assets_mime_type", "mime_type"),
CheckConstraint("size_bytes >= 0", name="ck_assets_size_nonneg"),
)
def __repr__(self) -> str:
return f"<Asset id={self.id} hash={(self.hash or '')[:12]}>"
class AssetReference(Base):
"""Unified model combining file cache state and user-facing metadata.
Each row represents either:
- A filesystem reference (file_path is set) with cache state
- An API-created reference (file_path is NULL) without cache state
"""
__tablename__ = "asset_references"
id: Mapped[str] = mapped_column(
String(36), primary_key=True, default=lambda: str(uuid.uuid4())
)
asset_id: Mapped[str] = mapped_column(
String(36), ForeignKey("assets.id", ondelete="CASCADE"), nullable=False
)
# Cache state fields (from former AssetCacheState)
file_path: Mapped[str | None] = mapped_column(Text, nullable=True)
# In-root loader path derived from file_path at scan/ingest time.
loader_path: Mapped[str | None] = mapped_column(Text, nullable=True)
mtime_ns: Mapped[int | None] = mapped_column(BigInteger, nullable=True)
needs_verify: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)
is_missing: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)
enrichment_level: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
# Info fields (from former AssetInfo)
owner_id: Mapped[str] = mapped_column(String(128), nullable=False, default="")
name: Mapped[str] = mapped_column(String(512), nullable=False)
mime_type: Mapped[str | None] = mapped_column(String(255))
system_metadata: Mapped[dict[str, Any] | None] = mapped_column(JSON(none_as_null=True))
job_id: Mapped[str | None] = mapped_column(String(36))
user_metadata: Mapped[dict[str, Any] | None] = mapped_column(JSON)
loader_path: Mapped[str | None] = mapped_column(Text)
preview_id: Mapped[str | None] = mapped_column(
String(36), ForeignKey("asset_references.id", ondelete="SET NULL")
String(36), ForeignKey("assets.id", ondelete="SET NULL")
)
user_metadata: Mapped[dict[str, Any] | None] = mapped_column(
JSON(none_as_null=True)
)
system_metadata: Mapped[dict[str, Any] | None] = mapped_column(
JSON(none_as_null=True), nullable=True, default=None
)
job_id: Mapped[str | None] = mapped_column(String(36), nullable=True, default=None)
created_at: Mapped[datetime] = mapped_column(
DateTime(timezone=False), nullable=False, default=get_utc_now
)
# Explicit user/API edit time; never use onupdate because row writes are not always edits.
updated_at: Mapped[datetime] = mapped_column(
DateTime(timezone=False), nullable=False, default=get_utc_now
)
last_access_time: Mapped[datetime] = mapped_column(
DateTime(timezone=False), nullable=False, default=get_utc_now
)
deleted_at: Mapped[datetime | None] = mapped_column(
DateTime(timezone=False), nullable=True, default=None
)
last_access_time: Mapped[datetime | None] = mapped_column(DateTime(timezone=False))
asset: Mapped[Asset] = relationship(
"Asset",
back_populates="references",
foreign_keys=[asset_id],
lazy="selectin",
content: Mapped[AssetContent] = relationship(back_populates="records", lazy="selectin")
preview: Mapped[Asset | None] = relationship(
"Asset", foreign_keys=[preview_id], remote_side=lambda: [Asset.id]
)
preview_ref: Mapped[AssetReference | None] = relationship(
"AssetReference",
foreign_keys=[preview_id],
remote_side=lambda: [AssetReference.id],
metadata_entries: Mapped[list[AssetMeta]] = relationship(
back_populates="asset", cascade="all,delete-orphan", passive_deletes=True
)
metadata_entries: Mapped[list[AssetReferenceMeta]] = relationship(
back_populates="asset_reference",
cascade="all,delete-orphan",
passive_deletes=True,
tag_links: Mapped[list[AssetTag]] = relationship(
back_populates="asset", cascade="all,delete-orphan", passive_deletes=True
)
tag_links: Mapped[list[AssetReferenceTag]] = relationship(
back_populates="asset_reference",
cascade="all,delete-orphan",
passive_deletes=True,
overlaps="tags,asset_references",
)
tags: Mapped[list[Tag]] = relationship(
secondary="asset_reference_tags",
back_populates="asset_references",
lazy="selectin",
viewonly=True,
overlaps="tag_links,asset_reference_links,asset_references,tag",
secondary="asset_tags", back_populates="assets", viewonly=True, lazy="selectin"
)
__table_args__ = (
Index("uq_asset_references_file_path", "file_path", unique=True),
Index("ix_asset_references_asset_id", "asset_id"),
Index("ix_asset_references_owner_id", "owner_id"),
Index("ix_asset_references_name", "name"),
Index("ix_asset_references_is_missing", "is_missing"),
Index("ix_asset_references_enrichment_level", "enrichment_level"),
Index("ix_asset_references_created_at", "created_at"),
Index("ix_asset_references_last_access_time", "last_access_time"),
Index("ix_asset_references_deleted_at", "deleted_at"),
Index("ix_asset_references_preview_id", "preview_id"),
Index("ix_asset_references_owner_name", "owner_id", "name"),
CheckConstraint(
"(mtime_ns IS NULL) OR (mtime_ns >= 0)", name="ck_ar_mtime_nonneg"
),
CheckConstraint(
"enrichment_level >= 0 AND enrichment_level <= 2",
name="ck_ar_enrichment_level_range",
),
Index("ix_assets_content_id", "content_id"),
Index("ix_assets_name", "name"),
Index("ix_assets_created_at", "created_at"),
Index("ix_assets_preview_id", "preview_id"),
)
def __repr__(self) -> str:
path_part = f" path={self.file_path!r}" if self.file_path else ""
return f"<AssetReference id={self.id} name={self.name!r}{path_part}>"
class AssetMeta(Base):
__tablename__ = "asset_meta"
class AssetReferenceMeta(Base):
__tablename__ = "asset_reference_meta"
asset_reference_id: Mapped[str] = mapped_column(
String(36),
ForeignKey("asset_references.id", ondelete="CASCADE"),
primary_key=True,
asset_id: Mapped[str] = mapped_column(
String(36), ForeignKey("assets.id", ondelete="CASCADE"), primary_key=True
)
key: Mapped[str] = mapped_column(String(256), primary_key=True)
ordinal: Mapped[int] = mapped_column(Integer, primary_key=True, default=0)
val_str: Mapped[str | None] = mapped_column(String(2048))
val_num: Mapped[Decimal | None] = mapped_column(Numeric(38, 10))
val_bool: Mapped[bool | None] = mapped_column(Boolean)
val_json: Mapped[Any | None] = mapped_column(JSON)
val_str: Mapped[str | None] = mapped_column(String(2048), nullable=True)
val_num: Mapped[float | None] = mapped_column(Numeric(38, 10), nullable=True)
val_bool: Mapped[bool | None] = mapped_column(Boolean, nullable=True)
val_json: Mapped[Any | None] = mapped_column(JSON(none_as_null=True), nullable=True)
asset_reference: Mapped[AssetReference] = relationship(
back_populates="metadata_entries"
)
asset: Mapped[Asset] = relationship(back_populates="metadata_entries")
__table_args__ = (
Index("ix_asset_reference_meta_key", "key"),
Index("ix_asset_reference_meta_key_val_str", "key", "val_str"),
Index("ix_asset_reference_meta_key_val_num", "key", "val_num"),
Index("ix_asset_reference_meta_key_val_bool", "key", "val_bool"),
Index("ix_asset_meta_key", "key"),
Index("ix_asset_meta_key_val_str", "key", "val_str"),
Index("ix_asset_meta_key_val_num", "key", "val_num"),
Index("ix_asset_meta_key_val_bool", "key", "val_bool"),
CheckConstraint(
"val_str IS NOT NULL OR val_num IS NOT NULL OR val_bool IS NOT NULL OR val_json IS NOT NULL",
name="has_value",
name="ck_asset_meta_has_value",
),
)
class AssetReferenceTag(Base):
__tablename__ = "asset_reference_tags"
class AssetTag(Base):
__tablename__ = "asset_tags"
asset_reference_id: Mapped[str] = mapped_column(
String(36),
ForeignKey("asset_references.id", ondelete="CASCADE"),
primary_key=True,
asset_id: Mapped[str] = mapped_column(
String(36), ForeignKey("assets.id", ondelete="CASCADE"), primary_key=True
)
tag_name: Mapped[str] = mapped_column(
String(512), ForeignKey("tags.name", ondelete="RESTRICT"), primary_key=True
@@ -216,12 +162,12 @@ class AssetReferenceTag(Base):
DateTime(timezone=False), nullable=False, default=get_utc_now
)
asset_reference: Mapped[AssetReference] = relationship(back_populates="tag_links")
tag: Mapped[Tag] = relationship(back_populates="asset_reference_links")
asset: Mapped[Asset] = relationship(back_populates="tag_links")
tag: Mapped[Tag] = relationship(back_populates="asset_links")
__table_args__ = (
Index("ix_asset_reference_tags_tag_name", "tag_name"),
Index("ix_asset_reference_tags_asset_reference_id", "asset_reference_id"),
Index("ix_asset_tags_tag_name", "tag_name"),
Index("ix_asset_tags_asset_id", "asset_id"),
)
@@ -229,17 +175,14 @@ class Tag(Base):
__tablename__ = "tags"
name: Mapped[str] = mapped_column(String(512), primary_key=True)
asset_reference_links: Mapped[list[AssetReferenceTag]] = relationship(
back_populates="tag",
overlaps="asset_references,tags",
)
asset_references: Mapped[list[AssetReference]] = relationship(
secondary="asset_reference_tags",
back_populates="tags",
viewonly=True,
overlaps="asset_reference_links,tag_links,tags,asset_reference",
asset_links: Mapped[list[AssetTag]] = relationship(back_populates="tag")
assets: Mapped[list[Asset]] = relationship(
secondary="asset_tags", back_populates="tags", viewonly=True
)
def __repr__(self) -> str:
return f"<Tag {self.name}>"
class AssetSystemState(Base):
__tablename__ = "asset_system_state"
key: Mapped[str] = mapped_column(String(256), primary_key=True)
value: Mapped[str] = mapped_column(Text, nullable=False)

View File

@@ -1,139 +1,44 @@
from app.assets.database.queries.asset import (
asset_exists_by_hash,
bulk_insert_assets,
create_stub_asset,
get_asset_by_hash,
get_existing_asset_ids,
reassign_asset_references,
update_asset_hash_and_mime,
upsert_asset,
)
from app.assets.database.queries.asset_reference import (
CacheStateRow,
UnenrichedReferenceRow,
bulk_insert_references_ignore_conflicts,
bulk_update_enrichment_level,
count_active_siblings,
bulk_update_is_missing,
bulk_update_needs_verify,
convert_metadata_to_rows,
delete_assets_by_ids,
delete_orphaned_seed_asset,
delete_reference_by_id,
delete_references_by_ids,
fetch_reference_and_asset,
fetch_reference_asset_and_tags,
get_or_create_reference,
get_reference_by_file_path,
get_reference_by_id,
get_reference_with_owner_check,
get_reference_ids_by_ids,
get_reference_paths_by_ids,
get_references_by_paths_and_asset_ids,
get_references_for_prefixes,
get_unenriched_references,
get_unreferenced_unhashed_asset_ids,
insert_reference,
list_all_file_paths_by_asset_id,
list_references_by_asset_id,
list_references_page,
mark_references_missing_outside_prefixes,
rebuild_metadata_projection,
reference_exists,
reference_exists_for_asset_id,
restore_references_by_paths,
set_reference_metadata,
set_reference_preview,
set_reference_system_metadata,
soft_delete_reference_by_id,
update_reference_access_time,
update_reference_name,
update_is_missing_by_asset_id,
update_reference_timestamps,
update_reference_updated_at,
upsert_reference,
)
from app.assets.database.queries.tags import (
AddTagsResult,
RemoveTagsResult,
SetTagsResult,
add_missing_tag_for_asset_id,
add_tags_to_reference,
bulk_insert_tags_and_meta,
ensure_tags_exist,
get_reference_tags,
list_tag_counts_for_filtered_assets,
list_tags_with_usage,
remove_missing_tag_for_asset_id,
remove_tags_from_reference,
set_reference_tags,
validate_tags_exist,
"""Re-exports the record and tag query functions so callers import them from one
place instead of reaching into individual query modules. A module-level
``__getattr__`` resolves names that live in the tag module, keeping this import
surface flat as queries are split across more files.
"""
from importlib import import_module
from app.assets.database.queries.records import (
create_content,
create_content_reporting_insert,
create_record,
delete_record,
fetch_record_tags,
get_record_by_id,
list_records_page,
mark_content_missing,
rename_record,
unset_content_missing,
update_record_access_time,
)
__all__ = [
"AddTagsResult",
"CacheStateRow",
"RemoveTagsResult",
"SetTagsResult",
"UnenrichedReferenceRow",
"add_missing_tag_for_asset_id",
"add_tags_to_reference",
"asset_exists_by_hash",
"bulk_insert_assets",
"bulk_insert_references_ignore_conflicts",
"bulk_insert_tags_and_meta",
"bulk_update_enrichment_level",
"count_active_siblings",
"create_stub_asset",
"bulk_update_is_missing",
"bulk_update_needs_verify",
"convert_metadata_to_rows",
"delete_assets_by_ids",
"delete_orphaned_seed_asset",
"delete_reference_by_id",
"delete_references_by_ids",
"ensure_tags_exist",
"fetch_reference_and_asset",
"fetch_reference_asset_and_tags",
"get_asset_by_hash",
"get_existing_asset_ids",
"get_or_create_reference",
"get_reference_by_file_path",
"get_reference_by_id",
"get_reference_with_owner_check",
"get_reference_ids_by_ids",
"get_reference_paths_by_ids",
"get_reference_tags",
"get_references_by_paths_and_asset_ids",
"get_references_for_prefixes",
"get_unenriched_references",
"get_unreferenced_unhashed_asset_ids",
"insert_reference",
"list_all_file_paths_by_asset_id",
"list_references_by_asset_id",
"list_references_page",
"list_tag_counts_for_filtered_assets",
"list_tags_with_usage",
"mark_references_missing_outside_prefixes",
"reassign_asset_references",
"rebuild_metadata_projection",
"reference_exists",
"reference_exists_for_asset_id",
"remove_missing_tag_for_asset_id",
"remove_tags_from_reference",
"restore_references_by_paths",
"set_reference_metadata",
"set_reference_preview",
"set_reference_system_metadata",
"soft_delete_reference_by_id",
"set_reference_tags",
"update_asset_hash_and_mime",
"update_is_missing_by_asset_id",
"update_reference_access_time",
"update_reference_name",
"update_reference_timestamps",
"update_reference_updated_at",
"upsert_asset",
"upsert_reference",
"validate_tags_exist",
"create_content",
"create_content_reporting_insert",
"create_record",
"delete_record",
"fetch_record_tags",
"get_record_by_id",
"list_records_page",
"mark_content_missing",
"rename_record",
"unset_content_missing",
"update_record_access_time",
]
def __getattr__(name: str):
for module_name in ("tags",):
module = import_module(f"app.assets.database.queries.{module_name}")
candidate = getattr(module, name, None)
if candidate is not None:
return candidate
raise AttributeError(name)

View File

@@ -1,152 +0,0 @@
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, AssetReference
from app.assets.database.queries.common import MAX_BIND_PARAMS, 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 not asset.mime_type:
asset.mime_type = mime_type
changed = True
if changed:
updated = True
return asset, created, updated
def create_stub_asset(
session: Session,
size_bytes: int,
mime_type: str | None = None,
) -> Asset:
"""Create a new asset with no hash (stub for later enrichment)."""
asset = Asset(size_bytes=size_bytes, mime_type=mime_type, hash=None)
session.add(asset)
session.flush()
return asset
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()
found: set[str] = set()
for chunk in iter_chunks(asset_ids, MAX_BIND_PARAMS):
rows = session.execute(
select(Asset.id).where(Asset.id.in_(chunk))
).fetchall()
found.update(row[0] for row in rows)
return found
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 and not asset.mime_type:
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.
"""
ref = session.get(AssetReference, reference_id)
if ref and ref.asset_id == from_asset_id:
ref.asset_id = to_asset_id
session.flush()

File diff suppressed because it is too large Load Diff

View File

@@ -1,138 +0,0 @@
"""Shared utilities for database query modules."""
import os
from decimal import Decimal
from typing import Iterable, Sequence
import sqlalchemy as sa
from sqlalchemy import exists
from app.assets.database.models import AssetReference, AssetReferenceMeta, AssetReferenceTag
from app.assets.helpers import escape_sql_like_string, normalize_tags
MAX_BIND_PARAMS = 800
def calculate_rows_per_statement(cols: int) -> int:
"""Calculate how many rows can fit in one statement given column count."""
return max(1, MAX_BIND_PARAMS // max(1, cols))
def iter_chunks(seq, n: int):
"""Yield successive n-sized chunks from seq."""
for i in range(0, len(seq), n):
yield seq[i : i + n]
def iter_row_chunks(rows: list[dict], cols_per_row: int) -> Iterable[list[dict]]:
"""Yield chunks of rows sized to fit within bind param limits."""
if not rows:
return
yield from iter_chunks(rows, calculate_rows_per_statement(cols_per_row))
def build_visible_owner_clause(owner_id: str) -> sa.sql.ClauseElement:
"""Build owner visibility predicate for reads.
Owner-less rows are visible to everyone.
"""
owner_id = (owner_id or "").strip()
if owner_id == "":
return AssetReference.owner_id == ""
return AssetReference.owner_id.in_(["", owner_id])
def build_prefix_like_conditions(
prefixes: list[str],
) -> list[sa.sql.ColumnElement]:
"""Build LIKE conditions for matching file paths under directory prefixes."""
conds = []
for p in prefixes:
base = os.path.abspath(p)
if not base.endswith(os.sep):
base += os.sep
escaped, esc = escape_sql_like_string(base)
conds.append(AssetReference.file_path.like(escaped + "%", escape=esc))
return conds
def apply_tag_filters(
stmt: sa.sql.Select,
include_tags: Sequence[str] | None = None,
exclude_tags: Sequence[str] | None = None,
any_tags: Sequence[str] | None = None,
) -> sa.sql.Select:
"""include_tags: every tag must be present; any_tags: at least one must be
present; exclude_tags: none may be present."""
include_tags = normalize_tags(include_tags)
exclude_tags = normalize_tags(exclude_tags)
any_tags = normalize_tags(any_tags)
if include_tags:
for tag_name in include_tags:
stmt = stmt.where(
exists().where(
(AssetReferenceTag.asset_reference_id == AssetReference.id)
& (AssetReferenceTag.tag_name == tag_name)
)
)
if any_tags:
stmt = stmt.where(
exists().where(
(AssetReferenceTag.asset_reference_id == AssetReference.id)
& (AssetReferenceTag.tag_name.in_(any_tags))
)
)
if exclude_tags:
stmt = stmt.where(
~exists().where(
(AssetReferenceTag.asset_reference_id == AssetReference.id)
& (AssetReferenceTag.tag_name.in_(exclude_tags))
)
)
return stmt
def apply_metadata_filter(
stmt: sa.sql.Select,
metadata_filter: dict | None = None,
) -> sa.sql.Select:
"""Apply filters using asset_reference_meta projection table."""
if not metadata_filter:
return stmt
def _exists_for_pred(key: str, *preds) -> sa.sql.ClauseElement:
return sa.exists().where(
AssetReferenceMeta.asset_reference_id == AssetReference.id,
AssetReferenceMeta.key == key,
*preds,
)
def _exists_clause_for_value(key: str, value) -> sa.sql.ClauseElement:
if value is None:
return sa.not_(
sa.exists().where(
AssetReferenceMeta.asset_reference_id == AssetReference.id,
AssetReferenceMeta.key == key,
)
)
if isinstance(value, bool):
return _exists_for_pred(key, AssetReferenceMeta.val_bool == bool(value))
if isinstance(value, (int, float, Decimal)):
num = value if isinstance(value, Decimal) else Decimal(str(value))
return _exists_for_pred(key, AssetReferenceMeta.val_num == num)
if isinstance(value, str):
return _exists_for_pred(key, AssetReferenceMeta.val_str == value)
return _exists_for_pred(key, AssetReferenceMeta.val_json == value)
for k, v in metadata_filter.items():
if isinstance(v, list):
ors = [_exists_clause_for_value(k, elem) for elem in v]
if ors:
stmt = stmt.where(sa.or_(*ors))
else:
stmt = stmt.where(_exists_clause_for_value(k, v))
return stmt

View File

@@ -0,0 +1,333 @@
"""Owns every write to content rows, records and their tag links, plus the paged
reads that list them. Inserts that can lose a race — a content row at a path, a
tag, a tag link — run inside a savepoint and re-read the conflicting row, so a
concurrent writer settles the call instead of raising, while a genuine
constraint failure still surfaces. This is the sole writer of a content row's
path, which is what lets other modules trust raw-column path predicates.
"""
from __future__ import annotations
import os
from collections.abc import Sequence
from datetime import datetime
from typing import Any, Literal, NamedTuple, TypeAlias
import sqlalchemy as sa
from sqlalchemy.exc import IntegrityError
from sqlalchemy.orm import Session, joinedload, noload
from sqlalchemy.sql.elements import ColumnElement
from app.assets.database.models import Asset, AssetContent, AssetTag, Tag
from app.assets.helpers import escape_sql_like_string, get_utc_now
RecordSortField: TypeAlias = Literal[
"name", "created_at", "updated_at", "size", "last_access_time"
]
RecordSortOrder: TypeAlias = Literal["asc", "desc"]
class RecordCursorBoundary(NamedTuple):
value: datetime | int | str
id: str
class RecordPageSpec(NamedTuple):
all_tags: tuple[str, ...] = ()
any_tags: tuple[str, ...] = ()
none_tags: tuple[str, ...] = ()
name_contains: str | None = None
limit: int = 20
offset: int = 0
sort: RecordSortField = "created_at"
order: RecordSortOrder = "desc"
after: RecordCursorBoundary | None = None
_LIVE_PATH_UNIQUE_INDEX = "uq_asset_contents_path_live"
def _is_live_path_conflict(error: IntegrityError) -> bool:
orig = error.orig
message = str(orig)
postgres_names_the_index = getattr(getattr(orig, "diag", None), "constraint_name", None) == _LIVE_PATH_UNIQUE_INDEX
sqlite_names_the_column = "UNIQUE constraint failed" in message and "asset_contents.path" in message
return postgres_names_the_index or sqlite_names_the_column
def create_content_reporting_insert(session: Session, path: str, hash: str | None = None, size_bytes: int = 0, mtime_ns: int | None = None) -> tuple[AssetContent, bool]:
# The sole writer of asset_contents.path, which is what makes the raw-column SQL prefix
# predicates sound — lifecycle's temp wipe HARD-DELETES every row its predicate admits.
path = os.path.abspath(path)
content = AssetContent(path=path, hash=hash, size_bytes=size_bytes, mtime_ns=mtime_ns)
try:
with session.begin_nested():
session.add(content)
session.flush()
return content, True
except IntegrityError as error:
if not _is_live_path_conflict(error):
raise
winner = session.execute(sa.select(AssetContent).where(AssetContent.path == path, AssetContent.is_missing.is_(False))).scalar_one()
return winner, False
def create_content(session: Session, path: str, hash: str | None = None, size_bytes: int = 0, mtime_ns: int | None = None) -> AssetContent:
content, _ = create_content_reporting_insert(session, path, hash, size_bytes, mtime_ns)
return content
def ensure_tag(session: Session, name: str) -> None:
if session.get(Tag, name) is not None:
return
try:
with session.begin_nested():
session.add(Tag(name=name))
session.flush()
except IntegrityError:
# Only a row that is now present proves a lost race; anything else is a real failure.
if session.get(Tag, name) is None:
raise
def ensure_tag_link(session: Session, *, asset_id: str, tag_name: str, origin: str) -> bool:
if session.get(AssetTag, (asset_id, tag_name)) is not None:
return False
try:
with session.begin_nested():
session.add(AssetTag(asset_id=asset_id, tag_name=tag_name, origin=origin))
session.flush()
except IntegrityError:
if session.get(AssetTag, (asset_id, tag_name)) is None:
raise
return False
return True
def create_record(session: Session, content_id: str, name: str, mime_type: str | None = None, job_id: str | None = None, loader_path: str | None = None, tags: Sequence[str] | None = None, *, system_metadata: dict[str, Any] | None = None) -> Asset:
record = Asset(content_id=content_id, name=name, mime_type=mime_type, job_id=job_id, loader_path=loader_path, system_metadata=system_metadata)
session.add(record)
session.flush()
for tag_name in dict.fromkeys(tags or ()):
ensure_tag(session, tag_name)
ensure_tag_link(session, asset_id=record.id, tag_name=tag_name, origin="manual")
session.flush()
return record
def get_record_by_id(session: Session, id: str) -> Asset | None:
return session.get(Asset, id)
def get_preview_file_paths_by_ids(
session: Session,
preview_ids: Sequence[str],
) -> dict[str, str]:
if not preview_ids:
return {}
rows = session.execute(
sa.select(Asset.id, AssetContent.path)
.join(AssetContent, Asset.content_id == AssetContent.id)
.where(
Asset.id.in_(preview_ids),
AssetContent.is_missing.is_(False),
)
)
return {preview_id: path for preview_id, path in rows}
def get_record_by_path_or_none(session: Session, path: str) -> Asset | None:
return session.scalar(
sa.select(Asset)
.join(AssetContent, Asset.content_id == AssetContent.id)
.where(AssetContent.path == path, AssetContent.is_missing.is_(False))
.order_by(Asset.created_at.desc(), Asset.id.desc())
.limit(1)
)
def fetch_record_tags(session: Session, record_id: str) -> list[str]:
return list(
session.scalars(
sa.select(AssetTag.tag_name)
.where(AssetTag.asset_id == record_id)
.order_by(AssetTag.tag_name)
)
)
def update_record_access_time(
session: Session,
record_id: str,
ts: datetime | None = None,
only_if_newer: bool = True,
) -> None:
ts = ts or get_utc_now()
stmt = sa.update(Asset).where(Asset.id == record_id)
if only_if_newer:
stmt = stmt.where(
sa.or_(
Asset.last_access_time.is_(None),
Asset.last_access_time < ts,
)
)
session.execute(stmt.values(last_access_time=ts))
def bump_record_updated_at(session: Session, record_id: str) -> None:
session.execute(
sa.update(Asset).where(Asset.id == record_id).values(updated_at=get_utc_now())
)
def build_record_tag_filter_clauses(
all_tags: Sequence[str],
any_tags: Sequence[str],
none_tags: Sequence[str],
) -> tuple[ColumnElement[bool], ...]:
clauses: list[ColumnElement[bool]] = []
for tag_name in all_tags:
clauses.append(
sa.exists(
sa.select(AssetTag.asset_id).where(
AssetTag.asset_id == Asset.id,
AssetTag.tag_name == tag_name,
)
)
)
if any_tags:
clauses.append(
sa.exists(
sa.select(AssetTag.asset_id).where(
AssetTag.asset_id == Asset.id,
AssetTag.tag_name.in_(any_tags),
)
)
)
if none_tags:
clauses.append(
~sa.exists(
sa.select(AssetTag.asset_id).where(
AssetTag.asset_id == Asset.id,
AssetTag.tag_name.in_(none_tags),
)
)
)
return tuple(clauses)
def list_records_page(
session: Session,
spec: RecordPageSpec,
) -> tuple[list[Asset], dict[str, list[str]], int]:
filters = list(build_record_tag_filter_clauses(spec.all_tags, spec.any_tags, spec.none_tags))
if spec.name_contains:
escaped_name, escape_character = escape_sql_like_string(spec.name_contains)
filters.append(
Asset.name.ilike(f"%{escaped_name}%", escape=escape_character)
)
sort_columns = {
"name": Asset.name,
"created_at": Asset.created_at,
"updated_at": Asset.updated_at,
"size": AssetContent.size_bytes,
"last_access_time": Asset.last_access_time,
}
sort_column = sort_columns[spec.sort]
descending = spec.order == "desc"
sort_expression = sort_column.desc() if descending else sort_column.asc()
id_expression = Asset.id.desc() if descending else Asset.id.asc()
statement = (
sa.select(Asset)
.join(AssetContent, Asset.content_id == AssetContent.id)
.where(*filters)
.options(joinedload(Asset.content), noload(Asset.tags))
)
if spec.after is not None:
comparison = (
sort_column < spec.after.value
if descending
else sort_column > spec.after.value
)
tied_comparison = (
Asset.id < spec.after.id
if descending
else Asset.id > spec.after.id
)
statement = statement.where(
sa.or_(
comparison,
sa.and_(
sort_column == spec.after.value,
tied_comparison,
),
)
)
statement = statement.order_by(sort_expression, id_expression).limit(spec.limit)
if spec.after is None:
statement = statement.offset(spec.offset)
records = list(session.scalars(statement))
total = session.scalar(
sa.select(sa.func.count())
.select_from(Asset)
.join(AssetContent, Asset.content_id == AssetContent.id)
.where(*filters)
)
record_ids = [record.id for record in records]
tag_map: dict[str, list[str]] = {}
if record_ids:
rows = session.execute(
sa.select(AssetTag.asset_id, AssetTag.tag_name)
.join(Asset, AssetTag.asset_id == Asset.id)
.join(AssetContent, Asset.content_id == AssetContent.id)
.where(AssetTag.asset_id.in_(record_ids))
.order_by(AssetTag.tag_name.asc())
)
for record_id, tag_name in rows:
tag_map.setdefault(record_id, []).append(tag_name)
return records, tag_map, int(total or 0)
def rename_record(session: Session, id: str, name: str) -> Asset:
record = session.get(Asset, id)
if record is None:
raise LookupError(id)
record.name = name
record.updated_at = get_utc_now()
session.flush()
return record
def delete_record(session: Session, id: str) -> None:
record = session.get(Asset, id)
if record is None:
return
session.delete(record)
session.flush()
def mark_content_missing(session: Session, content_id: str) -> None:
content = session.get(AssetContent, content_id)
if content is None:
raise LookupError(content_id)
content.is_missing = True
ensure_tag(session, "missing")
for record_id in session.scalars(sa.select(Asset.id).where(Asset.content_id == content_id)):
ensure_tag_link(session, asset_id=record_id, tag_name="missing", origin="automatic")
session.flush()
def unset_content_missing(session: Session, content_id: str) -> None:
content = session.get(AssetContent, content_id)
if content is None:
raise LookupError(content_id)
content.is_missing = False
session.execute(sa.delete(AssetTag).where(AssetTag.tag_name == "missing", AssetTag.asset_id.in_(sa.select(Asset.id).where(Asset.content_id == content_id))))
session.flush()

View File

@@ -1,26 +1,29 @@
from dataclasses import dataclass
from typing import Iterable, Sequence
"""Answers questions about tag usage: which tags exist with how many assets, and
how those counts narrow once a filter is applied. Both queries reuse the record
listing's own joins and filter clauses, so the counts a client sees always
describe the same assets the listing endpoint would return for that filter,
records with missing content included.
"""
import sqlalchemy as sa
from sqlalchemy import delete, func, select
from sqlalchemy.dialects import sqlite
from sqlalchemy.exc import IntegrityError
from __future__ import annotations
from dataclasses import dataclass, field
from typing import Sequence
from sqlalchemy import func, select
from sqlalchemy.orm import Session
from app.assets.database.models import (
Asset,
AssetReference,
AssetReferenceMeta,
AssetReferenceTag,
AssetContent,
AssetTag,
Tag,
)
from app.assets.database.queries.common import (
apply_metadata_filter,
apply_tag_filters,
build_visible_owner_clause,
iter_row_chunks,
from app.assets.database.queries.records import (
build_record_tag_filter_clauses,
)
from app.assets.helpers import escape_sql_like_string, get_utc_now, normalize_tags
from app.assets.helpers import escape_sql_like_string
@dataclass(frozen=True)
@@ -35,225 +38,7 @@ class RemoveTagsResult:
removed: list[str]
not_present: list[str]
total_tags: list[str]
@dataclass(frozen=True)
class SetTagsResult:
added: list[str]
removed: list[str]
total: list[str]
def validate_tags_exist(session: Session, tags: list[str]) -> None:
"""Raise ValueError if any of the given tag names do not exist."""
existing_tag_names = set(
name
for (name,) in session.execute(select(Tag.name).where(Tag.name.in_(tags))).all()
)
missing = [t for t in tags if t not in existing_tag_names]
if missing:
raise ValueError(f"Unknown tags: {missing}")
def ensure_tags_exist(session: Session, names: Iterable[str]) -> None:
wanted = normalize_tags(list(names))
if not wanted:
return
rows = [{"name": n} for n in list(dict.fromkeys(wanted))]
ins = (
sqlite.insert(Tag)
.values(rows)
.on_conflict_do_nothing(index_elements=[Tag.name])
)
session.execute(ins)
def get_reference_tags(session: Session, reference_id: str) -> list[str]:
return [
tag_name
for (tag_name,) in (
session.execute(
select(AssetReferenceTag.tag_name)
.where(AssetReferenceTag.asset_reference_id == reference_id)
.order_by(AssetReferenceTag.tag_name.asc())
)
).all()
]
def set_reference_tags(
session: Session,
reference_id: str,
tags: Sequence[str],
origin: str = "manual",
) -> SetTagsResult:
desired = normalize_tags(tags)
current = set(get_reference_tags(session, reference_id))
to_add = [t for t in desired if t not in current]
to_remove = [t for t in current if t not in desired]
if to_add:
ensure_tags_exist(session, to_add)
session.add_all(
[
AssetReferenceTag(
asset_reference_id=reference_id,
tag_name=t,
origin=origin,
added_at=get_utc_now(),
)
for t in to_add
]
)
session.flush()
if to_remove:
session.execute(
delete(AssetReferenceTag).where(
AssetReferenceTag.asset_reference_id == reference_id,
AssetReferenceTag.tag_name.in_(to_remove),
)
)
session.flush()
return SetTagsResult(added=sorted(to_add), removed=sorted(to_remove), total=sorted(desired))
def add_tags_to_reference(
session: Session,
reference_id: str,
tags: Sequence[str],
origin: str = "manual",
create_if_missing: bool = True,
reference_row: AssetReference | None = None,
) -> AddTagsResult:
if not reference_row:
ref = session.get(AssetReference, reference_id)
if not ref:
raise ValueError(f"AssetReference {reference_id} not found")
norm = normalize_tags(tags)
if not norm:
total = get_reference_tags(session, reference_id=reference_id)
return AddTagsResult(added=[], already_present=[], total_tags=total)
if create_if_missing:
ensure_tags_exist(session, norm)
current = set(get_reference_tags(session, reference_id))
want = set(norm)
to_add = sorted(want - current)
if to_add:
with session.begin_nested() as nested:
try:
session.add_all(
[
AssetReferenceTag(
asset_reference_id=reference_id,
tag_name=t,
origin=origin,
added_at=get_utc_now(),
)
for t in to_add
]
)
session.flush()
except IntegrityError:
nested.rollback()
after = set(get_reference_tags(session, reference_id=reference_id))
return AddTagsResult(
added=sorted(((after - current) & want)),
already_present=sorted(want & current),
total_tags=sorted(after),
)
def remove_tags_from_reference(
session: Session,
reference_id: str,
tags: Sequence[str],
) -> RemoveTagsResult:
ref = session.get(AssetReference, reference_id)
if not ref:
raise ValueError(f"AssetReference {reference_id} not found")
norm = normalize_tags(tags)
if not norm:
total = get_reference_tags(session, reference_id=reference_id)
return RemoveTagsResult(removed=[], not_present=[], total_tags=total)
existing = set(get_reference_tags(session, reference_id))
to_remove = sorted(set(t for t in norm if t in existing))
not_present = sorted(set(t for t in norm if t not in existing))
if to_remove:
session.execute(
delete(AssetReferenceTag).where(
AssetReferenceTag.asset_reference_id == reference_id,
AssetReferenceTag.tag_name.in_(to_remove),
)
)
session.flush()
total = get_reference_tags(session, reference_id=reference_id)
return RemoveTagsResult(removed=to_remove, not_present=not_present, total_tags=total)
def add_missing_tag_for_asset_id(
session: Session,
asset_id: str,
origin: str = "automatic",
) -> None:
select_rows = (
sa.select(
AssetReference.id.label("asset_reference_id"),
sa.literal("missing").label("tag_name"),
sa.literal(origin).label("origin"),
sa.literal(get_utc_now()).label("added_at"),
)
.where(AssetReference.asset_id == asset_id)
.where(
sa.not_(
sa.exists().where(
(AssetReferenceTag.asset_reference_id == AssetReference.id)
& (AssetReferenceTag.tag_name == "missing")
)
)
)
)
session.execute(
sqlite.insert(AssetReferenceTag)
.from_select(
["asset_reference_id", "tag_name", "origin", "added_at"],
select_rows,
)
.on_conflict_do_nothing(
index_elements=[
AssetReferenceTag.asset_reference_id,
AssetReferenceTag.tag_name,
]
)
)
def remove_missing_tag_for_asset_id(
session: Session,
asset_id: str,
) -> None:
session.execute(
sa.delete(AssetReferenceTag).where(
AssetReferenceTag.asset_reference_id.in_(
sa.select(AssetReference.id).where(AssetReference.asset_id == asset_id)
),
AssetReferenceTag.tag_name == "missing",
)
)
protected: list[str] = field(default_factory=list)
def list_tags_with_usage(
@@ -263,26 +48,18 @@ def list_tags_with_usage(
offset: int = 0,
include_zero: bool = True,
order: str = "count_desc",
owner_id: str = "",
) -> tuple[list[tuple[str, str, int]], int]:
) -> tuple[list[tuple[str, int]], int]:
prefix_filter = prefix.strip() if prefix else ""
counts_sq = (
select(
AssetReferenceTag.tag_name.label("tag_name"),
func.count(AssetReferenceTag.asset_reference_id).label("cnt"),
AssetTag.tag_name.label("tag_name"),
func.count(AssetTag.asset_id).label("cnt"),
)
.select_from(AssetReferenceTag)
.join(AssetReference, AssetReference.id == AssetReferenceTag.asset_reference_id)
.where(build_visible_owner_clause(owner_id))
.where(
sa.or_(
AssetReference.is_missing == False, # noqa: E712
AssetReferenceTag.tag_name == "missing",
)
)
.where(AssetReference.deleted_at.is_(None))
.group_by(AssetReferenceTag.tag_name)
.select_from(AssetTag)
.join(Asset, Asset.id == AssetTag.asset_id)
.join(AssetContent, Asset.content_id == AssetContent.id)
.group_by(AssetTag.tag_name)
.subquery()
)
@@ -311,17 +88,10 @@ def list_tags_with_usage(
total_q = total_q.where(func.substr(Tag.name, 1, len(prefix_filter)) == prefix_filter)
if not include_zero:
visible_tags_sq = (
select(AssetReferenceTag.tag_name)
.join(AssetReference, AssetReference.id == AssetReferenceTag.asset_reference_id)
.where(build_visible_owner_clause(owner_id))
.where(
sa.or_(
AssetReference.is_missing == False, # noqa: E712
AssetReferenceTag.tag_name == "missing",
)
)
.where(AssetReference.deleted_at.is_(None))
.group_by(AssetReferenceTag.tag_name)
select(AssetTag.tag_name)
.join(Asset, Asset.id == AssetTag.asset_id)
.join(AssetContent, Asset.content_id == AssetContent.id)
.group_by(AssetTag.tag_name)
)
total_q = total_q.where(Tag.name.in_(visible_tags_sq))
@@ -334,84 +104,48 @@ def list_tags_with_usage(
def list_tag_counts_for_filtered_assets(
session: Session,
owner_id: str = "",
include_tags: Sequence[str] | None = None,
exclude_tags: Sequence[str] | None = None,
name_contains: str | None = None,
metadata_filter: dict | None = None,
limit: int = 100,
# Appended last so pre-existing positional callers keep binding correctly.
any_tags: Sequence[str] | None = None,
) -> dict[str, int]:
"""Return tag counts for assets matching the given filters.
"""Return {tag_name: count} for the assets matching the given filters.
Uses the same filtering logic as list_references_page but returns
{tag_name: count} instead of paginated references.
Reuses build_record_tag_filter_clauses and the same Asset->AssetContent inner
join as list_records_page, so /api/assets and /api/assets/tags/refine agree on
which assets a given all/any/none + name_contains filter selects — including
missing-content records, which stay catalog-visible.
"""
# Build a subquery of matching reference IDs
ref_sq = (
select(AssetReference.id)
.join(Asset, Asset.id == AssetReference.asset_id)
.where(build_visible_owner_clause(owner_id))
.where(AssetReference.is_missing == False) # noqa: E712
.where(AssetReference.deleted_at.is_(None))
filters = list(
build_record_tag_filter_clauses(
tuple(include_tags or ()),
tuple(any_tags or ()),
tuple(exclude_tags or ()),
)
)
if name_contains:
escaped, esc = escape_sql_like_string(name_contains)
ref_sq = ref_sq.where(AssetReference.name.ilike(f"%{escaped}%", escape=esc))
filters.append(Asset.name.ilike(f"%{escaped}%", escape=esc))
ref_sq = apply_tag_filters(ref_sq, include_tags, exclude_tags, any_tags)
ref_sq = apply_metadata_filter(ref_sq, metadata_filter)
ref_sq = ref_sq.subquery()
asset_sq = (
select(Asset.id)
.join(AssetContent, Asset.content_id == AssetContent.id)
.where(*filters)
.subquery()
)
# Count tags across those references
q = (
select(
AssetReferenceTag.tag_name,
func.count(AssetReferenceTag.asset_reference_id).label("cnt"),
AssetTag.tag_name,
func.count(AssetTag.asset_id).label("cnt"),
)
.where(AssetReferenceTag.asset_reference_id.in_(select(ref_sq.c.id)))
.group_by(AssetReferenceTag.tag_name)
.order_by(func.count(AssetReferenceTag.asset_reference_id).desc(), AssetReferenceTag.tag_name.asc())
.where(AssetTag.asset_id.in_(select(asset_sq.c.id)))
.group_by(AssetTag.tag_name)
.order_by(func.count(AssetTag.asset_id).desc(), AssetTag.tag_name.asc())
.limit(limit)
)
rows = session.execute(q).all()
return {tag_name: int(cnt) for tag_name, cnt in rows}
def bulk_insert_tags_and_meta(
session: Session,
tag_rows: list[dict],
meta_rows: list[dict],
) -> None:
"""Batch insert into asset_reference_tags and asset_reference_meta.
Uses ON CONFLICT DO NOTHING.
Args:
session: Database session
tag_rows: Dicts with: asset_reference_id, tag_name, origin, added_at
meta_rows: Dicts with: asset_reference_id, key, ordinal, val_*
"""
if tag_rows:
ins_tags = sqlite.insert(AssetReferenceTag).on_conflict_do_nothing(
index_elements=[
AssetReferenceTag.asset_reference_id,
AssetReferenceTag.tag_name,
]
)
for chunk in iter_row_chunks(tag_rows, cols_per_row=4):
session.execute(ins_tags, chunk)
if meta_rows:
ins_meta = sqlite.insert(AssetReferenceMeta).on_conflict_do_nothing(
index_elements=[
AssetReferenceMeta.asset_reference_id,
AssetReferenceMeta.key,
AssetReferenceMeta.ordinal,
]
)
for chunk in iter_row_chunks(meta_rows, cols_per_row=7):
session.execute(ins_meta, chunk)

342
app/assets/docs/README.md Normal file
View File

@@ -0,0 +1,342 @@
# Assets
The asset system is only active when explicitly enabled at startup; it is off by default. When it is not enabled, asset API routes return a disabled-service error and no background filesystem scanning occurs. Routes answer that same disabled-service error when the system is enabled but its database dependencies are unavailable.
## Data model
An asset record represents one user-visible entity. It owns its name, tags, `job_id`, `loader_path`, MIME type, extracted metadata, user metadata, and optional preview relationship.
A content row represents bytes at a storage location. It owns the path, byte size, modification time, hash, and missing state. Multiple asset records may reference one content row. Two content rows may have the same hash; hash uniqueness is never a database invariant.
The path is required and unique among content rows that are not missing. Missing rows keep their last known path so the scanner can attempt recovery.
Names are labels and are not unique. Records are global; the local asset system has no per-user ownership boundary.
MIME type and extracted metadata belong to the asset record because the same bytes may be interpreted in different ways. Extracted metadata may grow when an extractor learns more, but an extractor must not silently remove facts it already recorded within the same extractor version.
`loader_path` is fixed when the record is created. It records how the registry classified the path at discovery time.
## Tags
Asset records carry two kinds of tags: user-applied tags added and removed through the tagging API, and system tags the backend derives on its own.
A location tag records which storage root a file sits under: `input`, `output`, or `temp` for files under those directories, and `models` for files under any configured model base directory. A file under a models base additionally gets one `model_type:<folder_name>` tag per model category whose base directory contains the file and whose registered extension set accepts the file's extension; a file whose extension matches no category still gets the `models` tag but no `model_type:` tag. Model-type tags come from the registered category names, not from path components.
On upload, tags choose the write destination. A request must carry exactly one destination role tag (`input`, `models`, or `output`), and a `models` upload must carry exactly one `model_type:<folder_name>` tag naming the category folder to write into. Any other tags in the request land on the created record as ordinary tags.
`missing` is the one system tag the tagging API refuses to add or remove; it is projected onto records whose content is absent, as described in Missing content. Location and model-type tags carry no such protection and can be added or removed like any other tag.
## Missing content
Every asset whose content is absent stays visible and carries the client-visible `missing` tag (see Tags). Missing state belongs to the content row, so all records that reference the same content become missing together. While content is missing, it cannot be downloaded or resolved by hash: attempts to fetch it or to look it up through from-hash resolution fail rather than silently succeeding.
A content row moves through these observable states:
```mermaid
stateDiagram-v2
LiveUnhashed: live, hash null
LiveHashed: live, hashed
MissingUnhashed: missing, hash null
MissingHashed: missing, hashed
[*] --> LiveUnhashed: scan discovery or output registration
[*] --> LiveHashed: upload (always hashed)
LiveUnhashed --> LiveHashed: hash computed (hashing on)
LiveHashed --> LiveUnhashed: mtime-only change (hashing off)
LiveUnhashed --> MissingUnhashed: file gone
LiveHashed --> MissingHashed: file gone
MissingHashed --> LiveHashed: hash match at path (hashing on)
MissingUnhashed --> LiveHashed: size and mtime match at path (hashing on)
```
A reappeared file that matches no missing candidate, or more than one, does not recover any of them. The scanner creates a new content row for the file instead, and the missing candidates stay missing. The same old-missing-plus-new-content shape applies to a same-path edit or reuse: the old row is marked missing and a separate new row and record are created for the new bytes, never transformed in place.
Only a hash can recover a missing content row, with one narrow exception for rows that were never hashed. The scanner hashes the file now present at the missing path and recovers the row only when exactly one missing candidate for that path has the same hash. If no candidate matches, the scanner creates new content. If multiple candidates match, none recover. Recovery also never fires for a path a live content row already occupies: the file there belongs to that row, and ordinary change detection owns it.
A missing row can have a null hash, for example a file deleted before it was ever hashed while every server was down during an off-to-on hashing transition. Such a row can never satisfy the hash comparison, so it has its own narrower recovery condition: it recovers only when it is the single missing null-hash candidate at that path, and its recorded byte size and modification time exactly match the reappeared file's verified stat. The freshly computed hash is set on the row as part of recovery. The stat facts are the only recorded identity a never-hashed row has, so a same-path reappearance that does not preserve them creates new content instead, and the old row stays missing rather than risk recovering the wrong record.
Hash-based recovery does not compare modification times; only the null-hash recovery described above requires the recorded byte size and modification time to match exactly. With hashing disabled, recovery is unavailable.
Missing rows and their asset records persist until explicitly deleted. Routine scans do not remove them.
## Hashing modes
The `--enable-asset-hashing` flag defaults to off. It controls whether the scanner and output pipeline hash file contents, as described below.
With hashing off, the scanner uses modification time and byte size to detect changes. A modification-time change with an unchanged byte size is treated as the same file: the stored file facts refresh and the stored hash is cleared, because without hashing the digest can no longer be vouched for. A change to both modification time and size means new content. A size change without a modification-time change is undefined behaviour.
With hashing on, the scanner uses modification time as the cheap change detector and a hash as the identity check. When the modification time changes, verification runs through the background seeding and enrichment work:
- If the hash is unchanged, refresh the stored file facts on the existing content row.
- If the hash changed, mark the old content missing and create a new content row and asset record for the path.
The system persists the previous hashing mode. On a transition from off to on, it hashes every live row without a hash and revalidates every live row that already has one. A file that cannot be read after a bounded number of attempts keeps its record live with its stored hash cleared, and the transition completes without it. On a transition from on to off, it retains existing hashes as inert data.
Hashing must use a stable file snapshot. The worker checks file facts before and after hashing and accepts the digest only when they still match each other. Otherwise it retries later.
### Uploads hash independently of the scanner flag
Upload dedup runs unconditionally in both hashing modes; only scanner and output hashing are gated by `--enable-asset-hashing`. Uploads always hash, use content-addressed filenames in hash-routed destinations, and deduplicate by bytes even while background hashing is off. Deduplication must know the digest before deciding whether to store the bytes, and the upload request is already reading the whole file.
From-hash lookup is disabled while hashing is off.
## Filesystem changes
### File removed outside the API
Mark the content missing and keep every asset record visible. The usual hash-based recovery applies if a file later appears at the same path.
### File edited in place
In hash mode, verify the new bytes. Refresh the existing content row if the hash is unchanged. If the hash differs, mark the old content missing and create new content and a new asset record for the path.
With hashing off, a change to both modification time and size is handled the same way as a hash difference: the old content is marked missing and new content and a new asset record are created. A modification-time change alone refreshes the existing row's file facts and clears its stored hash (see Hashing modes).
An asset record never changes from one byte identity to another. Existing history references therefore continue to point to the old, already-missing content instead of silently serving new bytes.
### Deleted path reused
Path reuse has the same final state as an in-place edit: the old content is missing and the new bytes receive new content and a new asset record. The result must not depend on whether a scan observed the deletion before the new file appeared.
### File moved or renamed outside the API
There is no filesystem move identity. Mark the old path missing and create new content and a new asset record at the new path.
### Partial download under its final filename
The scanner skips known partial-download extensions. For other files, it records file facts during the walk, waits once per scan pass for a short stability floor, and checks the facts again before inserting.
Files that change between those checks have not finished being written and are not admitted. Such a file is parked on a bounded, de-duplicated watch list. The watch list is rechecked more than once within a single scan pass: once during the fast walk phase and again during the enrichment phase. Each parked entry is retried a limited number of times before it is dropped; a dropped file is admitted normally whenever a later scan observes it in a stable state.
The watch list holds at most one entry per path and has a fixed maximum size. On overflow the oldest entry is evicted. Eviction is not data loss: an evicted path is rediscovered by ordinary directory traversal on any subsequent scan.
Retry is bounded by scan cadence, not by a timer. There is no background poller, and settlement is not guaranteed within any particular interval. A completed prompt queues an enrichment-only pass rather than a full filesystem walk, and that pass rechecks the watch list, so a partially-written output is normally admitted soon after the generation that produced it. If a stalled partial file passes the check and later resumes, normal change detection splits it from the prematurely admitted content.
### Symlinks and hardlinks
The scanner follows symlinks, stores lexical paths, and applies an inode cycle guard. Different paths produce different content rows even when they resolve to the same inode or bytes.
### Case and Unicode path forms
Store absolute, structurally normalised paths: relative segments and repeated separators collapse at the write boundary, so every stored path is the lexical absolute form. Never canonicalize case or Unicode; compare byte-for-byte. Filesystems that treat two case or Unicode spellings as the same path may therefore produce duplicate rows.
### Registry changes
Classification is fixed at record creation. A newly visible path receives the registry classification in effect when discovered. A path that leaves all registered prefixes becomes missing. An existing in-scope path is not reclassified when the registry changes.
## Asset operations
### Delete through the API
Delete the target asset record. Leave its content row and file intact. Do not soft-delete the record or revive the deleted identity during later discovery. The retained content row is not a tombstone: it is ordinary live content describing bytes that are still there, and nothing records that a deletion happened.
Deleting a record never deletes any other record. A preview record the deleted asset nominated stays untouched; references to a preview clear only when the preview record itself is deleted. The preview reference points from the deleted record to its target, so deleting the pointer must not destroy the target.
Content left behind after all its asset records are deleted can still be resolved by hash lookup, falling back to a generic name and a guessed content type when no record is left to supply one. There is currently no mechanism that reclaims or removes such orphaned content.
That retained row is also what makes the deletion durable. A scan seeds only paths with no live content row, so an unchanged file is never given a new record however often it is rediscovered. Only new content at that path produces a fresh asset record: bytes that changed and retired the old row, or content minted after the old row itself is gone. Such a record describes the new content and must never recreate the deleted identity.
### Rename
Renames always succeed. Duplicate names are allowed. Renaming does not change tags, classification, content identity, or storage path.
### Upload the same bytes with the same name
Every upload mints a new asset record for its own request. The request's tags, user metadata, and preview nomination always land on that new record, whether or not identical bytes were already uploaded under the same name. This applies to every upload endpoint, including `/upload/image`.
Content, not records, gets deduplicated. When qualifying content already holds these bytes (live, stat-consistent, and not temporary content), the new record points at that existing content row instead of the bytes being written again. When no content qualifies, the upload writes new content as normal. Either way, a new record is always created.
Uploads are not idempotent: a client that retries an upload after a timeout accumulates a second record rather than being handed back the first. Core has no idempotency protection anywhere, including prompt submission and deletion, and this matches the shape of cached output: a cache hit reuses content but still mints a new delivery record for the new caller.
Uploads hash in both hashing modes, so this content-level dedup applies regardless of the scanner hashing flag.
### `updated_at` semantics
An asset record's `updated_at` reflects only the last explicit user or API edit to that record: a rename, a user-metadata update, a MIME-type change, a preview nomination, or a manual tag add or remove.
It never advances for serving or downloading the asset (access time is a separate concern from edit time), for scanner enrichment filling in extracted metadata, for the automatic missing or recovered tag projection, for a content split or content retire, or for the preview-deleted foreign-key cascade that clears a `preview_id`. Reading or downloading an asset's content instead updates a separate last-access marker on the record or records involved; that marker only ever moves forward, never backward. Minting a new record, whether from an upload, a cached rerun, or a content split, does not touch any other record's `updated_at`. The new record carries its own fresh timestamp from creation, and every existing record's last-explicit-edit time stays untouched.
### Upload the same bytes with a different name
When byte matching is available, create a new asset record with the requested name and point it to the existing content row, wherever in the asset store that content lives. Do not write the bytes a second time. Without byte matching, create a new content row and write the upload normally.
The dedicated `/upload/image` endpoint scopes its byte-matching check to the same destination path only: it does not reuse content stored under a different path even when the hash matches, and writes a new copy there instead.
### Upload different bytes with the same name
Create a new asset record and new content. Both records keep the shared `name`. This is distinct from `display_name`, which is computed from the content's stored path and is not guaranteed to match `name` when content is stored under a hash-derived filename.
### Byte-identical generated output
Every non-cached save event creates a new asset record and a new content row. Do not merge generated outputs automatically. Once hashed, equal hashes make the byte relationship visible without changing either identity.
### Cached output
A fully cached rerun does not execute the save node or write a file. The asset layer creates a new asset record for the new prompt and points it to the existing content row for the cached file. It does not mutate the earlier asset record or content row.
The new record's extracted metadata is copied from the earliest existing record for that same content, ordered by creation, rather than re-extracted. When no earlier record exists for that content, metadata is extracted fresh.
If the cache is invalidated or unavailable, the save node executes normally. It writes a new counter-named file, and the same rule as ordinary generated output creates new content and a new asset record.
#### Runtime-expanded cached output
Whenever final history contains a runtime-expanded output locator, a child absent from the final executed-node set is treated exactly like any other cached output. It does not write a file. The asset layer creates a new delivery record that points to the existing content, and it does not mark earlier content or records missing or mutate them.
Execution must never be inferred from omission in a pre-execution cache announcement. A fully cached wrapper that produces no child locator creates no asset-registration event.
#### Registration happens at the point of the write
The asset layer registers a produced output at the moment the producing node finishes, not after the prompt completes. Classification is determined by control flow, never by comparing final history against a set of executed nodes:
- A node that reaches output processing has, by definition, executed. Cache hits return earlier and never reach it. Such an output is EXECUTED.
- An output delivered through the cached-UI path is, by definition, CACHED.
For an EXECUTED output the asset layer creates a new content row for the path, marks any existing live content row at that path missing, creates a new asset record carrying the current prompt's `job_id`, and returns that record's identifier to the caller.
For a CACHED output the asset layer creates a new asset record pointing at the existing content row with the new prompt's `job_id`, and mutates nothing else. The same holds for a fully cached wrapper that produces no child locator: no locator, no registration event.
Classifying or registering an output never requires a hash. Every non-cached save always creates new content regardless of whether the bytes changed, so classification never depends on comparing digests.
An identifier returned to a caller must refer to the record created for this write, never the identifier of a record that describes earlier bytes at the same path.
Registration failure must not fail a generation. The failure is logged, the output entry carries no asset identifier, and cleanup is attempted so that a partially-written asset row does not normally survive.
#### Emission and output identity
Both output paths register at the moment the output is emitted, the moment it becomes visible to anything outside the executor. Emission, not sending, is the boundary:
- The executed path emits when the producing node finishes processing its output.
- The cached path emits when a cached output is served.
At each emission the asset layer creates the appropriate record, and the resulting identifier is attached to the emitted output. Registration and publication to history occur unconditionally; only transmission to a connected client is conditional. An output must be registered whether or not a client is attached, so no registration may sit behind a client-connection check.
The database is the source of truth for output identity. For a fresh prompt, identity is always resolved through registration at emission, not by pulling a value back out of the cache. The cache mechanics used for subgraph replay can carry a previously-registered identifier forward as part of replaying that subgraph's own prior output; replay logic strips that identifier before treating the value as input to a fresh registration. A cached replay therefore reports the identifier of the delivery record created for the current prompt, never an identifier retained from the prompt that first produced the file.
It follows that no component reconstructs output classification after the fact by comparing final history against a record of which nodes executed. Classification is determined once, at emission, and that is the only source of truth for it.
### Settled properties and the eventually-consistent hash
Nearly every asset property that a client can observe is settled before that asset becomes observable, with two narrow exceptions described below: the content hash, and a scanner-discovered asset's metadata during its exception window.
A client reading an asset through the assets API never sees a property that is absent merely because work has not finished yet. Size, modification time, MIME type, extracted metadata, `job_id`, tags, preview location, and missing-state are all determined at the moment the record is created and are correct from the record's first observable instant.
The content hash is the sole eventually-consistent field. It is the only property whose cost scales with file size, so it is the only one permitted to be absent on a record a client can already see. A null hash means "not yet computed", never "this content has no hash".
Capabilities that require a hash (from-hash lookup, upload deduplication, and recovery of missing content) are unavailable for an asset whose hash has not yet settled. This is the same reduced capability described for disabled hashing. Settlement has no guaranteed deadline: a hash can remain unset indefinitely if the file never stabilizes or enrichment cannot make progress. It covers freshly-registered outputs even in hashing-on mode, not only the disabled-hashing case.
#### Exception: scanner-discovered assets
Assets discovered by the filesystem scanner are the one exception to full settlement: a seeded record can be observable before its metadata has settled.
Distinguishing "not yet enriched" from "enrichment ran and found nothing" matters here: without that distinction, excluding un-enriched records from view would hide, permanently, any asset whose extraction fails. Outputs and uploads are not covered by this exception; their properties are always settled when observable.
#### Output registration never hashes on the save path
Registering a produced output computes no content hash. The record is created with a null hash unconditionally, and the background enrichment pass fills it afterwards when hashing is enabled. When hashing is disabled the hash stays null, consistent with hashing being off.
Hashing reads every byte of the file and its cost scales without bound with file size, while the save path sits inside the execution loop. No output size may add hashing latency to a generation. No hashing failure may surface on the save path either; a failed or unstable read is the enrichment pass's problem, handled later.
Uploads are outside this behaviour and hash inline, always: deduplication must know the digest before deciding whether to store the bytes, and the upload request is already reading the whole file.
#### Two producers writing one path in a single prompt
Two output producers that resolve to the same path within one prompt are unsupported. The resulting database state is undefined and nothing here constrains it.
Registration occurs per producer at emission, so a prompt containing such a collision produces one record per producer, and their relative order determines the final state. That order is not a contract.
Each producer's registration attempt is independent and best-effort, consistent with the failure handling described for output registration: a failed attempt is caught and that producer's output proceeds without an asset identifier rather than failing the generation. Only the combined outcome across producers is undefined.
Save nodes assign counter-based filenames precisely so that concurrent producers do not collide, so a workflow reaching this state has bypassed the normal naming path.
### Server restart
Persist asset and content rows across restarts. In-memory history is transient. A record may retain a `job_id` whose prompt is no longer present in history.
In-memory tracking used during scanning, such as the unstable-file watch list and the hash-transition queue, does not persist across a restart either. A fresh process rebuilds whatever state it needs by re-observing the filesystem and database rather than resuming exactly where a prior process left off.
Temp records are the exception to row persistence (see Temp and preview output).
### Temp and preview output
Startup temp cleanup deletes both temp files and the records and content rows that represent them. Two failure modes are possible: if wiping the temp rows from the database fails, filesystem cleanup for that startup is skipped entirely and neither side is committed; if the database wipe succeeds but the subsequent filesystem removal fails, the rows are already gone from the database while the files remain on disk.
Cloud temp and preview assets expire through a separate mechanism.
### Temp content shared by permanent records
Expiry belongs to the content location. Upload dedup and from-hash lookup exclude temp content, so permanent records can never point to content that startup temp cleanup will remove.
### Same model bytes in two category locations
Create one asset record and one content row per location. Equal hashes may reveal that the bytes match. Do not merge the locations.
### Byte-identical content from two local users
The local asset system has global records and no owner field. It does not isolate or duplicate records by user.
### `/view` routes
`/view` accepts two query forms: a path-based form (`type`, `filename`, `subfolder`) and a blake3-hash form (`filename=blake3:<hash>`). `/api/assets/{id}/content` accepts an asset id and serves that asset's content directly. The blake3 form resolves only to non-temp content whose file is currently present (see Lookup and dedup against missing content).
### Lookup and dedup against missing content
Hash lookup, from-hash creation, and upload dedup consider only non-temp content whose file is present. A database row is not proof that bytes can be served.
When no qualifying content exists, uploads store the bytes they received and from-hash creation refuses. A missing sibling is never substituted for requested content.
### From-hash tie-breaking
When several qualifying content rows have the requested hash, choose the oldest by `created_at`, then by lexicographic `id`.
## Jobs and provenance
### Querying a job's outputs
The asset API does not provide a job-to-outputs query. A record's `job_id` is informational. In-memory history is the mechanism for showing a run's outputs during that session.
Cached reruns create a new asset record whose `job_id` is the new prompt. The record points to the existing content row because the save node did not execute and no new file was written. Earlier asset records keep their original `job_id` values.
### Recording where a file came from
Each asset record's `job_id` identifies the prompt associated with that record's own creation event. The system does not infer provenance across records that share content.
Embedded PNG metadata is independent of the asset database. Save code writes it into the file, and the frontend may read it to restore a workflow.
## Concurrency and failure boundaries
### File changes during hashing
Accept a digest only from a stable snapshot (see Hashing modes).
### Concurrent writers
Database constraints choose the winner when scanners, hooks, or uploads race within one process. Losing writers retry or discard their work. Tag and tag-link inserts run inside a savepoint and re-read the row they collided with, so a duplicate-key collision between concurrent tag writers settles quietly instead of surfacing an error. That covers duplicate keys only; lock contention is a separate limit (see Write pressure and reader starvation).
A file lock prevents more than one server process from opening the same database, so a second process never becomes a concurrent writer in the first place. Do not add advisory locks.
### Ambiguous recovery
If more than one missing content row matches a recreated file's hash, recover none of them.
### Database replacement while running
Deleting or replacing the database file while the server is running is undefined behaviour.
## Operational limits
### Write pressure and reader starvation
The asset database is SQLite with a single database-wide writer lock and no configured busy timeout or lock-error handling on any route. Several paths hold or contend for that lock: a non-deduplicated upload writes its bytes and mints a delivery record, while a deduplicated upload reuses existing content and mints only the record; a same-path write whose hash has changed retires the old content and inserts new content, while a same-path write whose hash matches refreshes the existing record in place; execution outputs register per-emission during the generation loop; a background enrichment pass fills hashes and metadata row by row; hash-serves write access time to every record sharing the served content; and the upload dedup claim holds the write lock across its filesystem re-check and metadata extraction.
Under sustained concurrent writes, a reader such as `GET /api/assets` can exceed SQLite's default five-second busy wait and surface an unhandled `database is locked` error as HTTP 500. The failure is transient and non-corrupting: no rows are corrupted, and a later request may succeed once the write pressure eases, though nothing retries or backs off automatically. No busy timeout, lock-error translation to 503, or WAL journal mode is configured.
## Schema migration
When startup finds the schema that predates the record/content split, it drops and recreates the affected tables inside the existing database file, then rebuilds them with a full scan. Rows are not migrated into the new schema, and the database file itself is not deleted. The database is backed up to a sibling file before the migration runs; if the migration fails, the database is reverted from that backup, and if it succeeds, the backup file is left in place rather than deleted.
Scanning runs in the background after startup returns, so the asset API can already be serving requests while the rebuilt tables are still being populated.
The rebuild discards data that a scan cannot reconstruct: manual tags, user metadata, preview assignments, API-created records, `job_id` links, any record renames, and deletions — nothing records which assets were deleted, so the scan mints a record for every file still on disk. When the rebuild runs, startup logs a warning naming the backup file and what was discarded.

View File

@@ -1,26 +1,45 @@
"""Small conversions the asset system needs in more than one layer: canonical
stored-hash strings, UTC timestamps, tag normalization, and a SQL predicate for
path containment. That predicate is deliberately case-sensitive and
component-bounded, because callers use it to choose rows for hard deletion.
"""
import os
from datetime import datetime, timezone
from typing import Sequence
import sqlalchemy as sa
from sqlalchemy.sql import ColumnElement
def select_best_live_path(states: Sequence) -> str:
def sql_path_under_prefix(
column: ColumnElement[str], prefix: str
) -> ColumnElement[bool]:
"""SQL predicate for ``Path(column).is_relative_to(prefix)`` on this platform.
Case-SENSITIVE and component-bounded: prefix ``/a/b`` matches ``/a/b`` and
``/a/b/c`` but not ``/a/bc``, ``/a/b-other`` or ``/a/B/c``.
``LIKE`` cannot express this. SQLite's ``LIKE`` is ASCII case-insensitive by
default, so ``'/data/TEMP/f' LIKE '/data/temp/%'`` is TRUE — which let the
temp wipe hard-delete records under a case-different persistent directory,
and let the enrichment scan mutate rows outside the requested root.
``GLOB`` is case-sensitive but carries its own metacharacters (``*``, ``?``,
``[``) with no ESCAPE clause, so every caller would need bracket-quoting.
``substr(column, 1, n) = <prefix>`` compares under the column's BINARY
collation and has no metacharacters at all, so a path containing ``%``,
``_``, ``*``, ``?`` or ``[`` needs no escaping and cannot inject.
Only the PREFIX is normalized here. That is sound because the column holds
normalized absolute paths — ``records.create_content`` is the sole writer
and normalizes there. Normalizing the column in SQL is not an option anyway:
it would need a per-row Python call and would defeat the index.
"""
Return the best on-disk path among cache states:
1) Prefer a path that exists with needs_verify == False (already verified).
2) Otherwise, pick the first path that exists.
3) Otherwise return empty string.
"""
alive = [
s
for s in states
if getattr(s, "file_path", None) and os.path.isfile(s.file_path)
]
if not alive:
return ""
for s in alive:
if not getattr(s, "needs_verify", False):
return s.file_path
return alive[0].file_path
base = os.path.abspath(prefix)
stem = base if base.endswith(os.sep) else base + os.sep
return sa.or_(
column == base,
sa.func.substr(column, 1, len(stem)) == stem,
)
def escape_sql_like_string(s: str, escape: str = "!") -> tuple[str, str]:
@@ -47,6 +66,10 @@ def normalize_tags(tags: list[str] | None) -> list[str]:
return list(dict.fromkeys(t.strip() for t in (tags or []) if (t or "").strip()))
def to_stored_hash(digest: str) -> str:
return f"blake3:{digest}"
def validate_blake3_hash(s: str) -> str:
"""Validate and normalize a blake3 hash string.

152
app/assets/lifecycle.py Normal file
View File

@@ -0,0 +1,152 @@
"""Runs the asset system's startup and shutdown maintenance: clearing temp rows
and files, recording a hash-mode transition, and handing the filesystem scan to
the background seeder. Startup only enqueues transition work — draining it
belongs to the seeder, so a large library cannot stall the server before it
accepts requests. It also settles, before any of that, whether the database
dependencies exist at all, disabling the asset routes when they do not.
"""
from __future__ import annotations
import logging
import os
import shutil
import folder_paths
from sqlalchemy import select
from app.assets.database.models import Asset, AssetContent
from app.assets.database.queries.records import delete_record
from app.assets.helpers import sql_path_under_prefix
from app.assets.services.hash_mode_state import enqueue_transition_work
from app.assets.services.hash_mode_state import record_transition_intent
from app.database.db import can_create_session, create_session
from comfy.cli_args import args
_excluded_scan_roots: set[str] = set()
_hash_mode_transition: str | None = None
def get_excluded_scan_roots() -> frozenset[str]:
return frozenset(_excluded_scan_roots)
def record_hash_mode_transition_intent() -> None:
global _hash_mode_transition
with create_session() as session:
_hash_mode_transition = record_transition_intent(session)
session.commit()
def enqueue_mode_transition_work() -> None:
with create_session() as session:
enqueue_transition_work(session, _hash_mode_transition)
session.commit()
def wipe_temp_db_rows(session) -> tuple[int, int]:
try:
temp_root = os.path.abspath(folder_paths.get_temp_directory())
except OSError:
return 0, 0
# These rows are hard-deleted, so the predicate must stay case-SENSITIVE: admitting a
# case-different persistent directory destroys user assets.
under_temp = sql_path_under_prefix(AssetContent.path, temp_root)
temp_record_ids = list(
session.scalars(
select(Asset.id)
.join(AssetContent, Asset.content_id == AssetContent.id)
.where(under_temp)
)
)
records_deleted = 0
for record_id in temp_record_ids:
delete_record(session, record_id)
records_deleted += 1
contents_deleted = 0
for content in session.scalars(select(AssetContent).where(under_temp)).all():
session.delete(content)
contents_deleted += 1
session.flush()
return records_deleted, contents_deleted
def cleanup_temp_filesystem() -> bool:
temp_dir = os.path.abspath(folder_paths.get_temp_directory())
if not os.path.exists(temp_dir):
return True
try:
shutil.rmtree(temp_dir)
return True
except OSError as exc:
logging.warning(
"Failed to remove temp directory %s: %s — excluding from scan for this process",
temp_dir,
exc,
)
_excluded_scan_roots.add(temp_dir)
return False
def start_asset_seeder() -> bool:
from app.assets.seeder import asset_seeder
started = asset_seeder.start(
roots=("models", "input", "output"),
prune_first=True,
compute_hashes=args.enable_asset_hashing,
)
if started:
logging.info("Background asset scan initiated for models, input, output")
return started
def run_asset_startup() -> None:
try:
with create_session() as session:
wipe_temp_db_rows(session)
session.commit()
except Exception:
logging.exception("Temp DB row wipe failed; skipping filesystem cleanup")
enqueue_mode_transition_work()
start_asset_seeder()
return
cleanup_temp_filesystem()
enqueue_mode_transition_work()
start_asset_seeder()
def run_startup(*, enable_assets: bool) -> None:
try:
if enable_assets:
run_asset_startup()
else:
cleanup_temp_filesystem()
except Exception:
logging.exception("Asset startup maintenance failed")
def run_asset_shutdown_cleanup() -> None:
try:
with create_session() as session:
wipe_temp_db_rows(session)
session.commit()
except Exception:
logging.exception("Temp DB row wipe failed during shutdown")
finally:
cleanup_temp_filesystem()
def run_shutdown() -> None:
try:
if can_create_session():
run_asset_shutdown_cleanup()
else:
cleanup_temp_filesystem()
except Exception:
logging.exception("Asset shutdown cleanup failed")

227
app/assets/manager.py Normal file
View File

@@ -0,0 +1,227 @@
import logging
from typing import Any, Callable, Protocol
from aiohttp import web
from app.assets import mode
from app.assets.api.routes import register_assets_routes
from app.assets.lifecycle import record_hash_mode_transition_intent, run_shutdown, run_startup
from app.assets.seeder import ScanPhase, asset_seeder
from app.assets.services.ingest import (
register_cached_output as ingest_register_cached_output,
register_executed_output as ingest_register_executed_output,
register_file_in_place,
)
from app.assets.services.path_utils import get_known_subfolder_tags
from app.assets.services.schemas import RegisteredAsset, UploadAssetView
from app.database.db import dependencies_available
from app.user_manager import UserManager
from comfy.cli_args import args
class AssetManager(Protocol):
@property
def enabled(self) -> bool: ...
def startup(self) -> None: ...
def shutdown(self) -> None: ...
def register_routes(
self, app: web.Application, user_manager: UserManager | None
) -> None: ...
def ensure_scan_started(self) -> None: ...
def pause_background_scan(self) -> None: ...
def queue_output_scan(self) -> None: ...
def resume_background_scan(self) -> None: ...
def register_upload(
self,
abs_path: str,
name: str,
upload_type: str,
subfolder: str,
*,
content_written: bool,
) -> UploadAssetView | None: ...
def register_executed_output(
self, abs_path: str, job_id: str | None
) -> RegisteredAsset | None: ...
def register_cached_output(
self, abs_path: str, job_id: str | None
) -> RegisteredAsset | None: ...
def set_event_sink(self, sink: Callable[[str, dict[str, Any]], None] | None) -> None: ...
class _ArgsLike(Protocol):
enable_assets: bool
enable_asset_hashing: bool
def _shutdown_assets() -> None:
asset_seeder.shutdown()
run_shutdown()
class NoAssets:
def __init__(self, args: _ArgsLike) -> None:
self._args = args
@property
def enabled(self) -> bool:
return False
def startup(self) -> None:
mode.init(self._args)
record_hash_mode_transition_intent()
run_startup(enable_assets=False)
def shutdown(self) -> None:
_shutdown_assets()
def register_routes(
self, app: web.Application, user_manager: UserManager | None
) -> None:
register_assets_routes(app)
asset_seeder.disable()
def ensure_scan_started(self) -> None:
return None
def pause_background_scan(self) -> None:
return None
def queue_output_scan(self) -> None:
return None
def resume_background_scan(self) -> None:
return None
def register_upload(
self,
abs_path: str,
name: str,
upload_type: str,
subfolder: str,
*,
content_written: bool,
) -> UploadAssetView | None:
return None
def register_executed_output(
self, abs_path: str, job_id: str | None
) -> RegisteredAsset | None:
return None
def register_cached_output(
self, abs_path: str, job_id: str | None
) -> RegisteredAsset | None:
return None
def set_event_sink(self, sink: Callable[[str, dict[str, Any]], None] | None) -> None:
return None
class AssetsEnabled:
def __init__(self, args: _ArgsLike) -> None:
self._args = args
@property
def enabled(self) -> bool:
return True
def startup(self) -> None:
mode.init(self._args)
record_hash_mode_transition_intent()
run_startup(enable_assets=True)
def shutdown(self) -> None:
_shutdown_assets()
def register_routes(
self, app: web.Application, user_manager: UserManager | None
) -> None:
register_assets_routes(app, user_manager)
def ensure_scan_started(self) -> None:
asset_seeder.start(roots=("models", "input", "output"))
def pause_background_scan(self) -> None:
asset_seeder.pause()
def queue_output_scan(self) -> None:
if not asset_seeder.is_disabled():
# FULL, not ENRICH: only a walk finds outputs a node never declared. Do not downgrade without re-weighing the cost.
asset_seeder.enqueue_scan(
roots=("output",),
phase=ScanPhase.FULL,
compute_hashes=self._args.enable_asset_hashing,
)
def resume_background_scan(self) -> None:
asset_seeder.resume()
def register_upload(
self,
abs_path: str,
name: str,
upload_type: str,
subfolder: str,
*,
content_written: bool,
) -> UploadAssetView | None:
try:
tag = upload_type if upload_type in ("input", "output") else "input"
tags = [tag] + get_known_subfolder_tags(subfolder)
result = register_file_in_place(
abs_path=abs_path,
name=name,
tags=tags,
content_written=content_written,
)
asset = RegisteredAsset(
id=result.ref.id,
content_id=result.content_id,
job_id=result.ref.job_id,
name=result.ref.name,
)
return UploadAssetView(
asset=asset,
asset_hash=result.asset.hash,
size=result.asset.size_bytes,
mime_type=result.asset.mime_type,
tags=result.tags,
)
except Exception:
logging.warning("Failed to register uploaded image as asset", exc_info=True)
return None
def register_executed_output(
self, abs_path: str, job_id: str | None
) -> RegisteredAsset | None:
return ingest_register_executed_output(abs_path, job_id)
def register_cached_output(
self, abs_path: str, job_id: str | None
) -> RegisteredAsset | None:
return ingest_register_cached_output(abs_path, job_id)
def set_event_sink(self, sink: Callable[[str, dict[str, Any]], None] | None) -> None:
asset_seeder.set_event_sink(sink)
def default_asset_manager() -> AssetManager:
if args.enable_assets and not dependencies_available():
logging.warning(
"Assets requested but database dependencies unavailable; asset endpoints "
"will answer 503. Please install the updated requirements.txt file."
)
return NoAssets(args)
return AssetsEnabled(args) if args.enable_assets else NoAssets(args)

30
app/assets/mode.py Normal file
View File

@@ -0,0 +1,30 @@
"""Holds whether asset hashing is enabled for this process, taken from the
command-line flag once at startup. Callers ask here rather than reading
arguments themselves, and asking before initialization raises instead of
defaulting, so a route can never quietly answer as though hashing were off.
"""
from __future__ import annotations
from typing import Protocol
class _HashingArguments(Protocol):
enable_asset_hashing: bool
_args: _HashingArguments | None = None
def init(args: _HashingArguments) -> None:
global _args
_args = args
def hashing_enabled() -> bool:
if _args is None:
raise RuntimeError(
"app.assets.mode.init() was not called before hashing_enabled(); "
"hash-mode state is uninitialised"
)
return bool(_args.enable_asset_hashing)

View File

@@ -1,66 +1,92 @@
"""Walks the asset roots and turns what it finds into catalog rows: collecting
paths, building specs, seeding new content and records, then enriching them
with metadata and hashes. Each spec is seeded inside its own savepoint, so one
file whose row conflicts cannot discard the work done for the files around it.
Enrichment counts as progress only when it produced what was asked of it — a
requested hash that could not be computed is no progress, which is what bounds
a pass over a file the server cannot read.
"""
import logging
import os
from dataclasses import dataclass
from pathlib import Path
from typing import Callable, Literal, TypedDict
import folder_paths
import sqlalchemy as sa
from sqlalchemy.exc import IntegrityError
from sqlalchemy.orm import Session
from app.assets import mode
from app.assets.database.queries import (
add_missing_tag_for_asset_id,
bulk_update_enrichment_level,
bulk_update_is_missing,
bulk_update_needs_verify,
delete_orphaned_seed_asset,
delete_references_by_ids,
ensure_tags_exist,
get_asset_by_hash,
get_reference_by_id,
get_references_for_prefixes,
get_unenriched_references,
mark_references_missing_outside_prefixes,
reassign_asset_references,
remove_missing_tag_for_asset_id,
set_reference_system_metadata,
update_asset_hash_and_mime,
create_content_reporting_insert,
mark_content_missing,
create_record,
)
from app.assets.services.bulk_ingest import (
SeedAssetSpec,
batch_insert_seed_assets,
from app.assets.database.models import Asset, AssetContent
from app.assets.helpers import sql_path_under_prefix, to_stored_hash
from app.assets.lifecycle import get_excluded_scan_roots
from app.assets.scanner_changes import (
clear_pending_verifications,
detect_content_change,
drain_pending_verifications,
is_path_under_prefixes,
live_contents_under_prefixes,
pending_recovery_count,
recover_missing_content,
)
from app.assets.services.file_utils import (
get_mtime_ns,
is_visible,
list_files_recursively,
verify_file_unchanged,
from app.assets.scanner_admission import (
PARTIAL_DOWNLOAD_EXTENSIONS as PARTIAL_DOWNLOAD_EXTENSIONS,
_WATCH_LIST as _WATCH_LIST,
_WatchEntry as _WatchEntry,
_should_skip_extension,
_two_stat_admit,
tick_watch_list as tick_watch_list,
)
from app.assets.services.hashing import HashCheckpoint, compute_blake3_hash
from app.assets.services.file_utils import get_mtime_ns, is_visible, list_files_recursively
from app.assets.services.image_dimensions import extract_image_dimensions
from app.assets.services.metadata_extract import extract_file_metadata
from app.assets.services.metadata_extract import ExtractedMetadata, extract_file_metadata
from app.assets.services.path_utils import (
compute_loader_path,
get_comfy_models_folders,
get_name_and_tags_from_asset_path,
)
from app.assets.services.ingest import _discard_unreferenced_content
from app.assets.services.snapshot_hash import snapshot_hash
from app.database.db import create_session
class _RefInfo(TypedDict):
ref_id: str
file_path: str
exists: bool
stat_unchanged: bool
needs_verify: bool
class _AssetAccumulator(TypedDict):
hash: str | None
size_db: int
refs: list[_RefInfo]
__all__ = [
"clear_pending_verifications",
"drain_pending_verifications",
"pending_recovery_count",
]
# Temp is deliberately absent: it is wiped before every scan, so walking it finds nothing.
RootType = Literal["models", "input", "output"]
class SeedAssetSpec(TypedDict):
abs_path: str
# Walk-time diagnostics only: seeding persists the seed-time restat instead.
size_bytes: int
mtime_ns: int
info_name: str
tags: list[str]
fname: str | None
metadata: ExtractedMetadata | None
mime_type: str | None
job_id: str | None
@dataclass(frozen=True, slots=True)
class UnenrichedContent:
content_id: str
record_id: str
file_path: str
def _log_scan_error(phase: str, error: OSError) -> None:
error_type = (
"permission_denied" if isinstance(error, PermissionError) else "os_error"
@@ -89,7 +115,10 @@ def get_owned_prefixes() -> list[str]:
def get_temp_prefixes() -> list[str]:
return [os.path.abspath(folder_paths.get_temp_directory())]
temp_dir = os.path.abspath(folder_paths.get_temp_directory())
if temp_dir in get_excluded_scan_roots():
return []
return [temp_dir]
def collect_models_files() -> list[str]:
@@ -118,147 +147,51 @@ def sync_references_with_filesystem(
session,
root: RootType,
collect_existing_paths: bool = False,
update_missing_tags: bool = False,
) -> set[str] | None:
return sync_prefixes_with_filesystem(
session,
get_scan_prefixes_for_root(root),
collect_existing_paths=collect_existing_paths,
update_missing_tags=update_missing_tags,
)
def sync_prefixes_with_filesystem(
session,
session: Session,
prefixes: list[str],
collect_existing_paths: bool = False,
update_missing_tags: bool = False,
) -> set[str] | None:
"""Reconcile asset references with filesystem under the given prefixes.
- Toggle needs_verify per reference using mtime/size stat check
- For hashed assets with at least one stat-unchanged ref: delete stale missing refs
- For seed assets with all refs missing: delete Asset and its references
- Optionally add/remove 'missing' tags based on stat check in this root
- Optionally return surviving absolute paths
Args:
session: Database session
prefixes: Absolute directory prefixes whose references to reconcile
collect_existing_paths: If True, return set of surviving file paths
update_missing_tags: If True, update 'missing' tags based on file status
Returns:
Set of surviving absolute paths if collect_existing_paths=True, else None
"""
if not prefixes:
return set() if collect_existing_paths else None
rows = get_references_for_prefixes(
session, prefixes, include_missing=update_missing_tags
)
by_asset: dict[str, _AssetAccumulator] = {}
for row in rows:
acc = by_asset.get(row.asset_id)
if acc is None:
acc = {"hash": row.asset_hash, "size_db": row.size_bytes, "refs": []}
by_asset[row.asset_id] = acc
stat_unchanged = False
try:
exists = True
stat_unchanged = verify_file_unchanged(
mtime_db=row.mtime_ns,
size_db=acc["size_db"],
stat_result=os.stat(row.file_path, follow_symlinks=True),
)
except FileNotFoundError:
exists = False
except PermissionError as e:
exists = True
_log_scan_error("reference_stat", e)
logging.debug("Permission denied accessing %s", row.file_path)
except OSError as e:
exists = False
_log_scan_error("reference_stat", e)
logging.debug("OSError checking %s: %s", row.file_path, e)
acc["refs"].append(
{
"ref_id": row.reference_id,
"file_path": row.file_path,
"exists": exists,
"stat_unchanged": stat_unchanged,
"needs_verify": row.needs_verify,
}
)
to_set_verify: list[str] = []
to_clear_verify: list[str] = []
stale_ref_ids: list[str] = []
to_mark_missing: list[str] = []
to_clear_missing: list[str] = []
survivors: set[str] = set()
for aid, acc in by_asset.items():
a_hash = acc["hash"]
refs = acc["refs"]
any_unchanged = any(r["stat_unchanged"] for r in refs)
all_missing = all(not r["exists"] for r in refs)
for r in refs:
if not r["exists"]:
to_mark_missing.append(r["ref_id"])
continue
if r["stat_unchanged"]:
to_clear_missing.append(r["ref_id"])
if r["needs_verify"]:
to_clear_verify.append(r["ref_id"])
if not r["stat_unchanged"] and not r["needs_verify"]:
to_set_verify.append(r["ref_id"])
if a_hash is None:
if refs and all_missing:
delete_orphaned_seed_asset(session, aid)
else:
for r in refs:
if r["exists"]:
survivors.add(os.path.abspath(r["file_path"]))
continue
if any_unchanged:
for r in refs:
if not r["exists"]:
stale_ref_ids.append(r["ref_id"])
if update_missing_tags:
try:
remove_missing_tag_for_asset_id(session, asset_id=aid)
except Exception as e:
logging.warning(
"Failed to remove missing tag for asset %s: %s", aid, e
)
elif update_missing_tags:
try:
add_missing_tag_for_asset_id(session, asset_id=aid, origin="automatic")
except Exception as e:
logging.warning("Failed to add missing tag for asset %s: %s", aid, e)
for r in refs:
if r["exists"]:
survivors.add(os.path.abspath(r["file_path"]))
delete_references_by_ids(session, stale_ref_ids)
stale_set = set(stale_ref_ids)
to_mark_missing = [ref_id for ref_id in to_mark_missing if ref_id not in stale_set]
bulk_update_is_missing(session, to_mark_missing, value=True)
bulk_update_is_missing(session, to_clear_missing, value=False)
bulk_update_needs_verify(session, to_set_verify, value=True)
bulk_update_needs_verify(session, to_clear_verify, value=False)
for content in live_contents_under_prefixes(session, prefixes):
try:
stat_result = os.stat(content.path, follow_symlinks=True)
except FileNotFoundError:
mark_content_missing(session, content.id)
except PermissionError as e:
_log_scan_error("reference_stat", e)
logging.debug("Permission denied accessing %s", content.path)
except OSError as e:
_log_scan_error("reference_stat", e)
logging.debug("OSError checking %s: %s", content.path, e)
mark_content_missing(session, content.id)
else:
detect_content_change(
session,
content,
stat_result,
hashing_is_enabled=mode.hashing_enabled(),
)
survivors.add(os.path.abspath(content.path))
return survivors if collect_existing_paths else None
def _is_under_prefixes(path: str, prefixes: list[str]) -> bool:
return is_path_under_prefixes(path, prefixes)
def sync_root_safely(root: RootType) -> set[str]:
"""Sync a single root's references with the filesystem.
@@ -270,7 +203,6 @@ def sync_root_safely(root: RootType) -> set[str]:
sess,
root,
collect_existing_paths=True,
update_missing_tags=True,
)
sess.commit()
return survivors or set()
@@ -296,7 +228,7 @@ def mark_missing_outside_prefixes_safely(prefixes: list[str]) -> int:
"""
try:
with create_session() as sess:
count = mark_references_missing_outside_prefixes(sess, prefixes)
count = mark_contents_missing_outside_prefixes(sess, prefixes)
sess.commit()
return count
except Exception as e:
@@ -304,6 +236,18 @@ def mark_missing_outside_prefixes_safely(prefixes: list[str]) -> int:
return 0
def mark_contents_missing_outside_prefixes(
session: Session, prefixes: list[str]
) -> int:
contents = session.scalars(
sa.select(AssetContent).where(AssetContent.is_missing.is_(False))
)
missing = [content for content in contents if not _is_under_prefixes(content.path, prefixes)]
for content in missing:
mark_content_missing(session, content.id)
return len(missing)
def collect_paths_for_roots(roots: tuple[RootType, ...]) -> list[str]:
"""Collect all file paths for the given roots."""
paths: list[str] = []
@@ -320,7 +264,6 @@ def build_asset_specs(
paths: list[str],
existing_paths: set[str],
enable_metadata_extraction: bool = True,
compute_hashes: bool = False,
) -> tuple[list[SeedAssetSpec], set[str], int]:
"""Build asset specs from paths, returning (specs, tag_pool, skipped_count).
@@ -328,14 +271,17 @@ def build_asset_specs(
paths: List of file paths to process
existing_paths: Set of paths that already exist in the database
enable_metadata_extraction: If True, extract tier 1 & 2 metadata
compute_hashes: If True, compute blake3 hashes (slow for large files)
"""
specs: list[SeedAssetSpec] = []
tag_pool: set[str] = set()
skipped = 0
candidates: list[tuple[str, os.stat_result]] = []
for p in paths:
abs_p = os.path.abspath(p)
if _should_skip_extension(abs_p):
skipped += 1
continue
if abs_p in existing_paths:
skipped += 1
continue
@@ -348,6 +294,12 @@ def build_asset_specs(
continue
if not stat_p.st_size:
continue
candidates.append((abs_p, stat_p))
admitted_paths, _ = _two_stat_admit(candidates)
candidate_stats = dict(candidates)
for abs_p in admitted_paths:
stat_p = candidate_stats[abs_p]
name, tags = get_name_and_tags_from_asset_path(abs_p)
rel_fname = compute_loader_path(abs_p)
@@ -360,15 +312,6 @@ def build_asset_specs(
relative_filename=rel_fname,
)
# Compute hash if requested
asset_hash: str | None = None
if compute_hashes:
try:
digest, _ = compute_blake3_hash(abs_p)
asset_hash = "blake3:" + digest
except Exception as e:
logging.warning("Failed to hash %s: %s", abs_p, e)
mime_type = metadata.content_type if metadata else None
specs.append(
{
@@ -379,7 +322,6 @@ def build_asset_specs(
"tags": tags,
"fname": rel_fname,
"metadata": metadata,
"hash": asset_hash,
"mime_type": mime_type,
"job_id": None,
}
@@ -389,40 +331,80 @@ def build_asset_specs(
return specs, tag_pool, skipped
def seed_asset_specs(session: Session, specs: list[SeedAssetSpec]) -> int:
created = 0
created_content_ids: list[str] = []
try:
for spec in specs:
path = os.path.abspath(spec["abs_path"])
try:
with session.begin_nested():
try:
stat_result = os.stat(path, follow_symlinks=True)
except OSError:
logging.warning("Skipping vanished asset during scan: %s", path)
continue
try:
recovery = recover_missing_content(
session,
path,
stat_result,
hashing_is_enabled=mode.hashing_enabled(),
)
except OSError:
logging.warning("Skipping vanished asset during scan: %s", path)
continue
if recovery != "no_match":
continue
content, inserted = create_content_reporting_insert(
session,
path=path,
hash=None,
size_bytes=stat_result.st_size,
mtime_ns=get_mtime_ns(stat_result),
)
if inserted:
created_content_ids.append(content.id)
existing_record = session.scalar(
sa.select(Asset.id).where(Asset.content_id == content.id).limit(1)
)
if existing_record is not None:
continue
create_record(
session,
content_id=content.id,
name=spec["info_name"],
mime_type=spec["mime_type"],
job_id=spec["job_id"],
loader_path=spec["fname"],
tags=spec["tags"],
)
created += 1
except IntegrityError:
logging.warning("Skipping asset whose row conflicts during scan: %s", path)
continue
except Exception:
session.rollback()
for content_id in created_content_ids:
_discard_unreferenced_content(session, content_id)
raise
return created
def insert_asset_specs(specs: list[SeedAssetSpec], tag_pool: set[str]) -> int:
"""Insert asset specs into database, returning count of created refs."""
def insert_asset_specs(specs: list[SeedAssetSpec], _tag_pool: set[str]) -> int:
if not specs:
return 0
with create_session() as sess:
if tag_pool:
ensure_tags_exist(sess, tag_pool)
result = batch_insert_seed_assets(sess, specs=specs, owner_id="")
created = seed_asset_specs(sess, specs)
sess.commit()
return result.inserted_refs
# Enrichment level constants
ENRICHMENT_STUB = 0 # Fast scan: path, size, mtime only
ENRICHMENT_METADATA = 1 # Metadata extracted (safetensors header, mime type)
ENRICHMENT_HASHED = 2 # Hash computed (blake3)
return created
def get_unenriched_assets_for_roots(
roots: tuple[RootType, ...],
max_level: int = ENRICHMENT_STUB,
compute_hashes: bool,
limit: int = 1000,
) -> list:
"""Get assets that need enrichment for the given roots.
Args:
roots: Tuple of root types to scan
max_level: Maximum enrichment level to include
limit: Maximum number of rows to return
Returns:
List of UnenrichedReferenceRow
"""
) -> list[UnenrichedContent]:
prefixes: list[str] = []
for root in roots:
prefixes.extend(get_scan_prefixes_for_root(root))
@@ -431,47 +413,61 @@ def get_unenriched_assets_for_roots(
return []
with create_session() as sess:
return get_unenriched_references(
sess, prefixes, max_level=max_level, limit=limit
query = (
sa.select(AssetContent.id, Asset.id, AssetContent.path)
.join(Asset, Asset.content_id == AssetContent.id)
.where(AssetContent.is_missing.is_(False))
)
if compute_hashes:
query = query.where(
sa.or_(
AssetContent.hash.is_(None),
Asset.system_metadata.is_(None),
)
)
else:
query = query.where(Asset.system_metadata.is_(None))
query = query.where(
sa.or_(
*(sql_path_under_prefix(AssetContent.path, p) for p in prefixes)
)
)
rows = sess.execute(query.order_by(Asset.id).limit(limit)).all()
return [
UnenrichedContent(content_id, record_id, file_path)
for content_id, record_id, file_path in rows
]
def enrich_asset(
session,
file_path: str,
reference_id: str,
asset_id: str,
content_id: str,
record_id: str,
extract_metadata: bool = True,
compute_hash: bool = False,
interrupt_check: Callable[[], bool] | None = None,
hash_checkpoints: dict[str, HashCheckpoint] | None = None,
) -> int:
) -> bool:
"""Enrich a single asset with metadata and/or hash.
Args:
session: Database session (caller manages lifecycle)
file_path: Absolute path to the file
reference_id: ID of the reference to update
asset_id: ID of the asset to update (for mime_type and hash)
content_id: ID of the content to update
record_id: ID of the record to update
extract_metadata: If True, extract safetensors header and mime type
compute_hash: If True, compute blake3 hash
interrupt_check: Optional non-blocking callable that returns True if
the operation should be interrupted (e.g. paused or cancelled)
hash_checkpoints: Optional dict for saving/restoring hash progress
across interruptions, keyed by file path
Returns:
New enrichment level achieved
Whether enrichment changed the B-schema record or content
"""
new_level = ENRICHMENT_STUB
try:
stat_p = os.stat(file_path, follow_symlinks=True)
except FileNotFoundError:
return new_level
return False
except OSError as e:
_log_scan_error("enrichment_stat", e)
return new_level
return False
initial_mtime_ns = get_mtime_ns(stat_p)
rel_fname = compute_loader_path(file_path)
@@ -486,71 +482,52 @@ def enrich_asset(
)
if metadata:
mime_type = metadata.content_type
new_level = ENRICHMENT_METADATA
full_hash: str | None = None
if compute_hash:
content = session.get(AssetContent, content_id)
digest: str | None = None
stored_hash: str | None = None
verified_stat: os.stat_result | None = None
hash_requested = compute_hash and content is not None and content.hash is None
if hash_requested:
try:
mtime_before = get_mtime_ns(stat_p)
size_before = stat_p.st_size
# Restore checkpoint if available and file unchanged
checkpoint = None
if hash_checkpoints is not None:
checkpoint = hash_checkpoints.get(file_path)
if checkpoint is not None:
cur_stat = os.stat(file_path, follow_symlinks=True)
if (checkpoint.mtime_ns != get_mtime_ns(cur_stat)
or checkpoint.file_size != cur_stat.st_size):
checkpoint = None
hash_checkpoints.pop(file_path, None)
else:
mtime_before = get_mtime_ns(cur_stat)
digest, new_checkpoint = compute_blake3_hash(
file_path,
interrupt_check=interrupt_check,
checkpoint=checkpoint,
)
if digest is None:
# Interrupted — save checkpoint for later resumption
if hash_checkpoints is not None and new_checkpoint is not None:
new_checkpoint.mtime_ns = mtime_before
new_checkpoint.file_size = size_before
hash_checkpoints[file_path] = new_checkpoint
return new_level
# Completed — clear any saved checkpoint
if hash_checkpoints is not None:
hash_checkpoints.pop(file_path, None)
stat_after = os.stat(file_path, follow_symlinks=True)
mtime_after = get_mtime_ns(stat_after)
if mtime_before != mtime_after:
logging.warning("File modified during hashing, discarding hash: %s", file_path)
else:
full_hash = f"blake3:{digest}"
metadata_ok = not extract_metadata or metadata is not None
if metadata_ok:
new_level = ENRICHMENT_HASHED
snapshot = snapshot_hash(file_path)
if snapshot is None:
logging.warning(
"File modified during hashing (snapshot unstable), discarding hash: %s",
file_path,
)
return False
digest, verified_stat = snapshot
stored_hash = to_stored_hash(digest)
except Exception as e:
if isinstance(e, OSError):
_log_scan_error("hashing", e)
else:
logging.warning("Failed to hash %s: %s", file_path, e)
# Optimistic guard: if the reference's mtime_ns changed since we
# started (e.g. ingest_existing_file updated it), our results are
# stale — discard them to avoid overwriting fresh registration data.
ref = get_reference_by_id(session, reference_id)
if ref is None or ref.mtime_ns != initial_mtime_ns:
record = session.get(Asset, record_id)
if content is None or record is None or content.mtime_ns != initial_mtime_ns:
session.rollback()
logging.info(
"Ref %s mtime changed during enrichment, discarding stale result",
reference_id,
"Content %s mtime changed during enrichment, discarding stale result",
content_id,
)
return ENRICHMENT_STUB
return False
# Non-NULL system_metadata permanently excludes the row from re-enrichment, so a
# disagreement here must discard the metadata too, not just the hash.
if verified_stat is not None and (
get_mtime_ns(verified_stat) != initial_mtime_ns
or verified_stat.st_size != stat_p.st_size
):
session.rollback()
logging.info(
"Content %s changed between its metadata read and its hash read, "
"discarding stale result",
content_id,
)
return False
if extract_metadata and metadata:
system_metadata = metadata.to_user_metadata()
@@ -558,32 +535,25 @@ def enrich_asset(
dims = extract_image_dimensions(file_path, mime_type=mime_type)
if dims:
system_metadata.update(dims)
set_reference_system_metadata(session, reference_id, system_metadata)
record.system_metadata = {**(record.system_metadata or {}), **system_metadata}
if full_hash:
existing = get_asset_by_hash(session, full_hash)
if existing and existing.id != asset_id:
reassign_asset_references(session, asset_id, existing.id, reference_id)
delete_orphaned_seed_asset(session, asset_id)
if mime_type:
update_asset_hash_and_mime(session, existing.id, mime_type=mime_type)
else:
update_asset_hash_and_mime(session, asset_id, full_hash, mime_type)
elif mime_type:
update_asset_hash_and_mime(session, asset_id, mime_type=mime_type)
if stored_hash:
content.hash = stored_hash
if mime_type:
record.mime_type = mime_type
bulk_update_enrichment_level(session, [reference_id], new_level)
session.commit()
return new_level
if hash_requested and stored_hash is None:
return False
return stored_hash is not None or metadata is not None or mime_type is not None
def enrich_assets_batch(
rows: list,
rows: list[UnenrichedContent],
extract_metadata: bool = True,
compute_hash: bool = False,
interrupt_check: Callable[[], bool] | None = None,
hash_checkpoints: dict[str, HashCheckpoint] | None = None,
) -> tuple[int, list[str]]:
"""Enrich a batch of assets.
@@ -597,8 +567,6 @@ def enrich_assets_batch(
compute_hash: If True, compute hash for each asset
interrupt_check: Optional non-blocking callable that returns True if
the operation should be interrupted (e.g. paused or cancelled)
hash_checkpoints: Optional dict for saving/restoring hash progress
across interruptions, keyed by file path
Returns:
Tuple of (enriched_count, failed_reference_ids)
@@ -612,23 +580,21 @@ def enrich_assets_batch(
break
try:
new_level = enrich_asset(
updated = enrich_asset(
sess,
file_path=row.file_path,
reference_id=row.reference_id,
asset_id=row.asset_id,
content_id=row.content_id,
record_id=row.record_id,
extract_metadata=extract_metadata,
compute_hash=compute_hash,
interrupt_check=interrupt_check,
hash_checkpoints=hash_checkpoints,
)
if new_level > row.enrichment_level:
if updated:
enriched += 1
else:
failed_ids.append(row.reference_id)
failed_ids.append(row.record_id)
except Exception as e:
logging.warning("Failed to enrich %s: %s", row.file_path, e)
sess.rollback()
failed_ids.append(row.reference_id)
failed_ids.append(row.record_id)
return enriched, failed_ids

View File

@@ -0,0 +1,97 @@
"""Decides whether a file observed during a filesystem walk is settled enough to
catalog, and parks the ones that are not. Admission requires two stats taken a
moment apart to agree; a file still changing, or carrying a partial-download
extension, goes onto a bounded watch list that later scans re-check until it
settles, vanishes, or exhausts its retries. This is what keeps a model that is
still downloading out of the catalog.
"""
from __future__ import annotations
import mimetypes
import os
import time
from dataclasses import dataclass
from typing import Final
from sqlalchemy.orm import Session
from app.assets.services.path_utils import compute_loader_path, get_name_and_tags_from_asset_path
PARTIAL_DOWNLOAD_EXTENSIONS = frozenset({
".part", ".partial", ".crdownload", ".download", ".tmp", ".aria2", ".!qb", ".opdownload",
})
_WATCH_SCAN_RETRIES: Final = 30
_WATCH_LIST_MAX_SIZE: Final = 256
@dataclass
class _WatchEntry:
path: str
last_stat: os.stat_result
ticks: int = 0
_WATCH_LIST: list[_WatchEntry] = []
def _should_skip_extension(path: str) -> bool:
return os.path.splitext(path)[1].lower() in PARTIAL_DOWNLOAD_EXTENSIONS
def _two_stat_admit(paths_with_stats: list[tuple[str, os.stat_result]]) -> tuple[list[str], list[str]]:
if not paths_with_stats:
return [], []
time.sleep(0.1)
admitted: list[str] = []
watched: list[str] = []
for path, first_stat in paths_with_stats:
try:
second_stat = os.stat(path)
except FileNotFoundError:
continue
if (second_stat.st_mtime_ns, second_stat.st_size) == (first_stat.st_mtime_ns, first_stat.st_size):
_WATCH_LIST[:] = [entry for entry in _WATCH_LIST if entry.path != path]
admitted.append(path)
else:
for entry in _WATCH_LIST:
if entry.path == path:
entry.last_stat = second_stat
break
else:
_WATCH_LIST.append(_WatchEntry(path, second_stat))
if len(_WATCH_LIST) > _WATCH_LIST_MAX_SIZE:
_ = _WATCH_LIST.pop(0)
watched.append(path)
return admitted, watched
def tick_watch_list(session: Session) -> None:
from app.assets.scanner import seed_asset_specs, SeedAssetSpec
remaining: list[_WatchEntry] = []
for entry in _WATCH_LIST:
try:
current = os.stat(entry.path)
except FileNotFoundError:
continue
if (current.st_mtime_ns, current.st_size) == (entry.last_stat.st_mtime_ns, entry.last_stat.st_size):
name, tags = get_name_and_tags_from_asset_path(entry.path)
spec: SeedAssetSpec = {
"abs_path": entry.path,
"size_bytes": current.st_size,
"mtime_ns": current.st_mtime_ns,
"info_name": name,
"tags": tags,
"fname": compute_loader_path(entry.path),
"metadata": None,
"mime_type": mimetypes.guess_type(entry.path, strict=False)[0],
"job_id": None,
}
seed_asset_specs(session, [spec])
continue
entry.last_stat = current
entry.ticks += 1
if entry.ticks < _WATCH_SCAN_RETRIES:
remaining.append(entry)
_WATCH_LIST[:] = remaining

View File

@@ -0,0 +1,205 @@
"""Reconciles catalogued content against what is actually on disk: retiring rows
whose file is gone, splitting a row whose bytes changed, and recovering one
whose file came back. Recovery fires only when the returning file's hash
identifies exactly one missing row and no live row already occupies that path,
so a restored file can never leave two live rows describing one location.
"""
from __future__ import annotations
import os
from pathlib import Path
from typing import Literal
import sqlalchemy as sa
from sqlalchemy.orm import Session
from app.assets.database.models import AssetContent
from app.assets.database.queries.records import (
create_content,
create_record,
mark_content_missing,
unset_content_missing,
)
from app.assets.helpers import sql_path_under_prefix, to_stored_hash
from app.assets.services.path_utils import compute_loader_path, get_name_and_tags_from_asset_path
from app.assets.services.snapshot_hash import snapshot_hash
_pending_verification_ids: list[str] = []
_pending_recovery_paths: list[str] = []
def clear_pending_verifications() -> None:
_pending_verification_ids.clear()
_pending_recovery_paths.clear()
def queue_pending_verification(content_id: str) -> None:
if content_id not in _pending_verification_ids:
_pending_verification_ids.append(content_id)
def pending_recovery_count() -> int:
return len(_pending_recovery_paths)
def recover_missing_content(
session: Session, path: str, stat_result: os.stat_result, hashing_is_enabled: bool
) -> Literal["recovered", "no_match", "unstable"]:
if not hashing_is_enabled:
return "no_match"
occupied = session.scalar(
sa.select(AssetContent.id)
.where(AssetContent.path == path, AssetContent.is_missing.is_(False))
.limit(1)
)
if occupied is not None:
return "no_match"
snapshot = snapshot_hash(path)
if snapshot is None:
if path not in _pending_recovery_paths:
_pending_recovery_paths.append(path)
return "unstable"
digest, verified_stat = snapshot
stored_hash = to_stored_hash(digest)
matches = list(
session.scalars(
sa.select(AssetContent).where(
AssetContent.path == path,
AssetContent.is_missing.is_(True),
AssetContent.hash == stored_hash,
)
)
)
if len(matches) == 1:
recovered = matches[0]
unset_content_missing(session, recovered.id)
recovered.size_bytes = verified_stat.st_size
recovered.mtime_ns = verified_stat.st_mtime_ns
return "recovered"
if len(matches) > 1:
return "no_match"
null_hash_matches = list(
session.scalars(
sa.select(AssetContent).where(
AssetContent.path == path,
AssetContent.is_missing.is_(True),
AssetContent.hash.is_(None),
)
)
)
if len(null_hash_matches) != 1:
return "no_match"
candidate = null_hash_matches[0]
if (candidate.size_bytes, candidate.mtime_ns) != (
verified_stat.st_size,
verified_stat.st_mtime_ns,
):
return "no_match"
unset_content_missing(session, candidate.id)
candidate.hash = stored_hash
candidate.size_bytes = verified_stat.st_size
candidate.mtime_ns = verified_stat.st_mtime_ns
return "recovered"
def is_path_under_prefixes(path: str, prefixes: list[str]) -> bool:
candidate = Path(os.path.abspath(path))
return any(candidate.is_relative_to(os.path.abspath(prefix)) for prefix in prefixes)
def split_content(session: Session, content: AssetContent, stat_result: os.stat_result, hash_value: str | None) -> AssetContent:
mark_content_missing(session, content.id)
name, tags = get_name_and_tags_from_asset_path(content.path)
replacement = create_content(
session,
path=content.path,
hash=hash_value,
size_bytes=stat_result.st_size,
mtime_ns=stat_result.st_mtime_ns,
)
create_record(
session,
content_id=replacement.id,
name=name,
loader_path=compute_loader_path(content.path),
tags=tags,
)
return replacement
def detect_content_change(
session: Session,
content: AssetContent,
stat_result: os.stat_result,
hashing_is_enabled: bool,
) -> None:
if content.mtime_ns == stat_result.st_mtime_ns:
# Ruling #10: size drift with unchanged mtime is undefined behavior.
return
if hashing_is_enabled:
queue_pending_verification(content.id)
return
if content.size_bytes == stat_result.st_size:
# User identity rule: a same-size mtime bump (rsync, cloud sync, backup restore) is the
# same file — never split, or the record's tags and metadata are destroyed.
# The stored hash goes with the refreshed stat: OFF mode cannot prove the bytes, and a
# refreshed stat alone would re-qualify the row to be served under a digest it may no
# longer match.
content.size_bytes = stat_result.st_size
content.mtime_ns = stat_result.st_mtime_ns
content.hash = None
return
split_content(session, content, stat_result, hash_value=None)
def drain_pending_verifications(session: Session, limit: int | None = None) -> int:
queued_count = min(len(_pending_verification_ids), limit or len(_pending_verification_ids))
processed = 0
for _ in range(queued_count):
content_id = _pending_verification_ids.pop(0)
content = session.get(AssetContent, content_id)
if content is None or content.is_missing:
continue
try:
os.stat(content.path, follow_symlinks=True)
except FileNotFoundError:
mark_content_missing(session, content.id)
processed += 1
continue
except OSError:
queue_pending_verification(content_id)
continue
try:
snapshot = snapshot_hash(content.path)
except OSError:
queue_pending_verification(content_id)
continue
if snapshot is None:
queue_pending_verification(content_id)
continue
digest, verified_stat = snapshot
stored_hash = to_stored_hash(digest)
if content.hash == stored_hash or content.hash is None:
content.hash = stored_hash
content.size_bytes = verified_stat.st_size
content.mtime_ns = verified_stat.st_mtime_ns
else:
split_content(session, content, verified_stat, hash_value=stored_hash)
processed += 1
return processed
def live_contents_under_prefixes(session: Session, prefixes: list[str]) -> list[AssetContent]:
if not prefixes:
return []
return list(
session.scalars(
sa.select(AssetContent).where(
AssetContent.is_missing.is_(False),
sa.or_(*(sql_path_under_prefix(AssetContent.path, prefix) for prefix in prefixes)),
)
)
)

View File

@@ -1,4 +1,10 @@
"""Background asset seeder with thread management and cancellation support."""
"""Runs the filesystem scan on a background thread so startup never waits for it,
exposing pause, resume, cancel and progress to the API. A run seeds
newly-observed files first, then enriches records in batches, and settles any
pending hash-mode transition at the start of the enrich phase so a server that
receives no prompts still completes the switch. A pass stops once batches stop
making progress, bounding a scan over files that cannot be read.
"""
import logging
import os
@@ -6,11 +12,9 @@ import threading
import time
from dataclasses import dataclass, field
from enum import Enum
from typing import Callable
from typing import Any, Callable, TypedDict
from app.assets.scanner import (
ENRICHMENT_METADATA,
ENRICHMENT_STUB,
RootType,
build_asset_specs,
collect_paths_for_roots,
@@ -22,8 +26,11 @@ from app.assets.scanner import (
mark_missing_outside_prefixes_safely,
sync_root_safely,
sync_temp_references_safely,
drain_pending_verifications,
tick_watch_list,
)
from app.database.db import dependencies_available
from app.assets.services.hash_mode_state import drain_transition_queue, pending_transition_count
from app.database.db import create_session, dependencies_available
class ScanInProgressError(Exception):
@@ -47,6 +54,12 @@ class ScanPhase(Enum):
FULL = "full" # Both phases sequentially
class PendingScan(TypedDict):
roots: tuple[RootType, ...]
phase: ScanPhase
compute_hashes: bool
@dataclass
class Progress:
"""Progress information for a scan operation."""
@@ -94,8 +107,12 @@ class _AssetSeeder:
self._compute_hashes: bool = False
self._prune_first: bool = False
self._progress_callback: ProgressCallback | None = None
self._event_sink: Callable[[str, dict[str, Any]], None] | None = None
self._disabled: bool = False
self._pending_enrich: dict | None = None
self._pending_scan: PendingScan | None = None
def set_event_sink(self, sink: Callable[[str, dict[str, Any]], None] | None) -> None:
self._event_sink = sink
def disable(self) -> None:
"""Disable the asset seeder, preventing any scans from starting."""
@@ -113,6 +130,8 @@ class _AssetSeeder:
progress_callback: ProgressCallback | None = None,
prune_first: bool = False,
compute_hashes: bool = False,
*,
_start_paused: bool = False,
) -> bool:
"""Start a background scan for the given roots.
@@ -122,6 +141,7 @@ class _AssetSeeder:
progress_callback: Optional callback called with progress updates
prune_first: If True, prune orphaned assets before scanning
compute_hashes: If True, compute blake3 hashes (slow)
_start_paused: Start with phase work blocked until resume()
Returns:
True if scan was started, False if already running
@@ -134,7 +154,7 @@ class _AssetSeeder:
if self._state != State.IDLE:
logging.info("Asset seeder already running, skipping start")
return False
self._state = State.RUNNING
self._state = State.PAUSED if _start_paused else State.RUNNING
self._progress = Progress()
self._errors = []
self._roots = roots
@@ -143,7 +163,10 @@ class _AssetSeeder:
self._compute_hashes = compute_hashes
self._progress_callback = progress_callback
self._cancel_event.clear()
self._run_gate.set() # Ensure unpaused when starting
if _start_paused:
self._run_gate.clear()
else:
self._run_gate.set()
self._thread = threading.Thread(
target=self._run_scan,
name="_AssetSeeder",
@@ -176,64 +199,40 @@ class _AssetSeeder:
compute_hashes=False,
)
def start_enrich(
def enqueue_scan(
self,
roots: tuple[RootType, ...] = ("models", "input", "output"),
progress_callback: ProgressCallback | None = None,
roots: tuple[RootType, ...],
phase: ScanPhase,
compute_hashes: bool = False,
) -> bool:
"""Start an enrichment scan (phase 2 only) - extracts metadata and hashes.
Args:
roots: Tuple of root types to scan
progress_callback: Optional callback for progress updates
compute_hashes: If True, compute blake3 hashes
Returns:
True if scan was started, False if already running
"""
return self.start(
roots=roots,
phase=ScanPhase.ENRICH,
progress_callback=progress_callback,
prune_first=False,
compute_hashes=compute_hashes,
)
def enqueue_enrich(
self,
roots: tuple[RootType, ...] = ("models", "input", "output"),
compute_hashes: bool = False,
) -> bool:
"""Start an enrichment scan now, or queue it for after the current scan.
If the seeder is idle, starts immediately. Otherwise, the enrich
request is stored and will run automatically when the current scan
finishes.
Args:
roots: Tuple of root types to scan
compute_hashes: If True, compute blake3 hashes
Returns:
True if started immediately, False if queued for later
"""
with self._lock:
if self.start_enrich(roots=roots, compute_hashes=compute_hashes):
if self.start(
roots=roots,
phase=phase,
prune_first=False,
compute_hashes=compute_hashes,
):
return True
if self._pending_enrich is not None:
existing_roots = set(self._pending_enrich["roots"])
if self._pending_scan is not None:
existing_roots = set(self._pending_scan["roots"])
existing_roots.update(roots)
self._pending_enrich["roots"] = tuple(existing_roots)
self._pending_enrich["compute_hashes"] = (
self._pending_enrich["compute_hashes"] or compute_hashes
self._pending_scan["roots"] = tuple(existing_roots)
self._pending_scan["compute_hashes"] = (
self._pending_scan["compute_hashes"] or compute_hashes
)
if self._pending_scan["phase"] is not phase:
self._pending_scan["phase"] = ScanPhase.FULL
else:
self._pending_enrich = {
self._pending_scan = {
"roots": roots,
"phase": phase,
"compute_hashes": compute_hashes,
}
logging.info("Enrich scan queued (roots=%s)", self._pending_enrich["roots"])
logging.info(
"Scan queued (roots=%s, phase=%s)",
self._pending_scan["roots"],
self._pending_scan["phase"].value,
)
return False
def cancel(self) -> bool:
@@ -370,16 +369,26 @@ class _AssetSeeder:
errors=list(self._errors),
)
def shutdown(self, timeout: float = 5.0) -> None:
def shutdown(self, timeout: float = 5.0) -> bool:
"""Gracefully shutdown: cancel any running scan and wait for thread.
Args:
timeout: Maximum seconds to wait for thread to exit
Returns:
True if the scan thread joined cleanly; False on timeout.
"""
self.cancel()
self.wait(timeout=timeout)
joined = self.wait(timeout=timeout)
if not joined:
logging.warning(
"Asset seeder thread did not exit within %ss",
timeout,
)
with self._lock:
self._thread = None
if joined:
self._thread = None
return joined
def mark_missing_outside_prefixes(self) -> int:
"""Mark references as missing when outside all known root prefixes.
@@ -458,13 +467,11 @@ class _AssetSeeder:
self._run_gate.wait() # Blocks if paused
return self._is_cancelled()
def _emit_event(self, event_type: str, data: dict) -> None:
def _emit_event(self, event_type: str, data: dict[str, Any]) -> None:
"""Emit a WebSocket event if server is available."""
try:
from server import PromptServer
if hasattr(PromptServer, "instance") and PromptServer.instance:
PromptServer.instance.send_sync(event_type, data)
if self._event_sink is not None:
self._event_sink(event_type, data)
except Exception:
pass
@@ -639,17 +646,22 @@ class _AssetSeeder:
},
)
with self._lock:
start_paused = self._state is State.PAUSED
self._reset_to_idle()
pending = self._pending_enrich
pending = self._pending_scan
if pending is not None:
self._pending_enrich = None
if not self.start_enrich(
self._pending_scan = None
if not self.start(
roots=pending["roots"],
phase=pending["phase"],
prune_first=False,
compute_hashes=pending["compute_hashes"],
_start_paused=start_paused,
):
logging.warning(
"Pending enrich scan could not start (roots=%s)",
"Pending scan could not start (roots=%s, phase=%s)",
pending["roots"],
pending["phase"].value,
)
def _run_fast_phase(self, roots: tuple[RootType, ...]) -> tuple[int, int, int]:
@@ -698,7 +710,6 @@ class _AssetSeeder:
paths,
existing_paths,
enable_metadata_extraction=False,
compute_hashes=False,
)
logging.debug(
"Fast scan: build_asset_specs took %.3fs (%d specs, %d skipped)",
@@ -751,6 +762,9 @@ class _AssetSeeder:
last_progress_time = now
self._update_progress(scanned=len(specs), created=total_created)
with create_session() as session:
tick_watch_list(session)
session.commit()
logging.info(
"Fast scan complete: %.3fs total (created=%d, skipped=%d, total_paths=%d)",
time.perf_counter() - t_fast_start,
@@ -767,16 +781,18 @@ class _AssetSeeder:
Tuple of (cancelled, total_enriched)
"""
total_enriched = 0
with create_session() as session:
drain_pending_verifications(session)
tick_watch_list(session)
for _ in range(3):
drain_transition_queue(session)
session.commit()
if pending_transition_count() == 0:
break
batch_size = 100
last_progress_time = time.perf_counter()
progress_interval = 1.0
# Get the target enrichment level based on compute_hashes
if not self._compute_hashes:
target_max_level = ENRICHMENT_STUB
else:
target_max_level = ENRICHMENT_METADATA
self._emit_event(
"assets.seed.started",
{"roots": list(roots), "phase": "enrich"},
@@ -786,10 +802,6 @@ class _AssetSeeder:
consecutive_empty = 0
max_consecutive_empty = 3
# Hash checkpoints survive across batches so interrupted hashes
# can be resumed without re-reading the entire file.
hash_checkpoints: dict[str, object] = {}
while True:
if self._check_pause_and_cancel():
logging.info("Enrich scan cancelled after %d assets", total_enriched)
@@ -798,13 +810,13 @@ class _AssetSeeder:
# Fetch next batch of unenriched assets
unenriched = get_unenriched_assets_for_roots(
roots,
max_level=target_max_level,
compute_hashes=self._compute_hashes,
limit=batch_size,
)
# Filter out previously failed references
if skip_ids:
unenriched = [r for r in unenriched if r.reference_id not in skip_ids]
unenriched = [row for row in unenriched if row.record_id not in skip_ids]
if not unenriched:
break
@@ -814,7 +826,6 @@ class _AssetSeeder:
extract_metadata=True,
compute_hash=self._compute_hashes,
interrupt_check=self._is_paused_or_cancelled,
hash_checkpoints=hash_checkpoints,
)
total_enriched += enriched
skip_ids.update(failed_ids)

View File

@@ -1,93 +1,43 @@
from app.assets.services.asset_management import (
asset_exists,
delete_asset_reference,
get_asset_by_hash,
get_asset_detail,
list_assets_page,
get_preview_file_paths,
resolve_asset_for_download,
set_asset_preview,
update_asset_metadata,
)
from app.assets.services.bulk_ingest import (
BulkInsertResult,
batch_insert_seed_assets,
cleanup_unreferenced_assets,
)
from app.assets.services.file_utils import (
get_mtime_ns,
get_size_and_mtime_ns,
list_files_recursively,
verify_file_unchanged,
)
"""Re-exports the asset service functions so routes and the seeder import them
from one place instead of reaching into individual service modules.
"""
from app.assets.services.ingest import (
DependencyMissingError,
HashMismatchError,
create_from_hash,
ingest_existing_file,
register_output_files,
UploadUnstableError,
upload_from_temp_path,
create_from_hash,
register_file_in_place,
)
from app.assets.database.queries import (
AddTagsResult,
RemoveTagsResult,
)
from app.assets.services.schemas import (
AssetData,
AssetDetailResult,
AssetSummaryData,
DownloadResolutionResult,
IngestResult,
ListAssetsResult,
ReferenceData,
RegisterAssetResult,
TagUsage,
UploadResult,
UserMetadata,
from app.assets.services.asset_management import (
get_asset_detail,
update_asset_metadata,
delete_asset_reference,
asset_exists,
get_preview_file_paths,
resolve_asset_for_download,
)
from app.assets.services.tagging import (
apply_tags,
list_tags,
remove_tags,
list_tags,
)
__all__ = [
"AddTagsResult",
"AssetData",
"AssetDetailResult",
"AssetSummaryData",
"ReferenceData",
"BulkInsertResult",
"DependencyMissingError",
"DownloadResolutionResult",
"HashMismatchError",
"IngestResult",
"ListAssetsResult",
"RegisterAssetResult",
"RemoveTagsResult",
"TagUsage",
"UploadResult",
"UserMetadata",
"apply_tags",
"asset_exists",
"batch_insert_seed_assets",
"UploadUnstableError",
"upload_from_temp_path",
"create_from_hash",
"delete_asset_reference",
"get_asset_by_hash",
"register_file_in_place",
"get_asset_detail",
"ingest_existing_file",
"register_output_files",
"get_mtime_ns",
"get_size_and_mtime_ns",
"list_assets_page",
"list_files_recursively",
"list_tags",
"cleanup_unreferenced_assets",
"remove_tags",
"update_asset_metadata",
"delete_asset_reference",
"asset_exists",
"get_preview_file_paths",
"resolve_asset_for_download",
"set_asset_preview",
"update_asset_metadata",
"upload_from_temp_path",
"verify_file_unchanged",
"apply_tags",
"remove_tags",
"list_tags",
]

View File

@@ -1,77 +1,87 @@
import contextlib
"""Serves the per-asset operations behind the API: reading an asset's detail,
updating its name, tags, metadata and preview, deleting a record, and resolving
a hash to a servable path. An update moves ``updated_at`` only when it actually
changed something, so a call that requests what is already true is not recorded
as a user edit.
"""
import mimetypes
import os
from datetime import timezone
from typing import Sequence
from app.assets.services.cursor import (
CursorPayload,
InvalidCursorError,
decode_cursor,
decode_cursor_int,
decode_cursor_time,
encode_cursor,
encode_cursor_from_time,
)
from sqlalchemy import delete, select, update
from app.assets.database.models import Asset
from app.assets.database.models import Asset, AssetContent, AssetTag
from app.assets.database.queries import (
asset_exists_by_hash,
reference_exists_for_asset_id,
delete_reference_by_id,
fetch_reference_and_asset,
get_reference_paths_by_ids,
soft_delete_reference_by_id,
fetch_reference_asset_and_tags,
get_asset_by_hash as queries_get_asset_by_hash,
get_reference_by_id,
get_reference_with_owner_check,
list_references_page,
list_all_file_paths_by_asset_id,
list_references_by_asset_id,
set_reference_metadata,
set_reference_preview,
set_reference_tags,
update_asset_hash_and_mime,
update_reference_access_time,
update_reference_name,
update_reference_updated_at,
delete_record,
fetch_record_tags,
get_record_by_id,
update_record_access_time,
)
from app.assets.helpers import select_best_live_path
from app.assets.services.path_utils import compute_loader_path
from app.assets.database.queries.records import (
bump_record_updated_at,
ensure_tag,
ensure_tag_link,
get_preview_file_paths_by_ids,
rename_record,
)
from app.assets.helpers import get_utc_now, normalize_tags, validate_blake3_hash
from app.assets.services.lookup import lookup_for_view
from app.assets.services.schemas import (
AssetData,
AssetDetailResult,
AssetSummaryData,
DownloadResolutionResult,
ListAssetsResult,
ReferenceData,
UserMetadata,
extract_asset_data,
extract_reference_data,
)
from app.database.db import create_session
def _record_to_detail_result(session, record) -> AssetDetailResult:
content = session.get(AssetContent, record.content_id)
tags = fetch_record_tags(session, record.id)
api_hash = content.hash if content else None
ref = ReferenceData(
id=record.id,
name=record.name,
file_path=content.path if content else None,
loader_path=record.loader_path,
user_metadata=record.user_metadata,
preview_id=record.preview_id,
system_metadata=record.system_metadata,
job_id=record.job_id,
created_at=record.created_at,
updated_at=record.updated_at,
last_access_time=record.last_access_time,
)
asset = AssetData(
hash=api_hash,
size_bytes=content.size_bytes if content else None,
mime_type=record.mime_type,
is_missing=bool(content and content.is_missing),
)
return AssetDetailResult(ref=ref, asset=asset, tags=tags)
def get_asset_detail(
reference_id: str,
owner_id: str = "",
) -> AssetDetailResult | None:
with create_session() as session:
result = fetch_reference_asset_and_tags(
session,
reference_id=reference_id,
owner_id=owner_id,
)
if not result:
record = get_record_by_id(session, reference_id)
if record is None:
return None
return _record_to_detail_result(session, record)
ref, asset, tags = result
return AssetDetailResult(
ref=extract_reference_data(ref),
asset=extract_asset_data(asset),
tags=tags,
def _fetch_manual_tags(session, reference_id: str) -> set[str]:
return set(
session.scalars(
select(AssetTag.tag_name).where(
AssetTag.asset_id == reference_id,
AssetTag.origin != "automatic",
)
)
)
def update_asset_metadata(
@@ -80,342 +90,132 @@ def update_asset_metadata(
tags: Sequence[str] | None = None,
user_metadata: UserMetadata = None,
tag_origin: str = "manual",
owner_id: str = "",
mime_type: str | None = None,
preview_id: str | None = None,
) -> AssetDetailResult:
with create_session() as session:
ref = get_reference_with_owner_check(session, reference_id, owner_id)
record = get_record_by_id(session, reference_id)
if record is None:
raise ValueError(f"Asset {reference_id} not found")
touched = False
if name is not None and name != ref.name:
update_reference_name(session, reference_id=reference_id, name=name)
touched = True
computed_filename = compute_loader_path(ref.file_path) if ref.file_path else None
new_meta: dict | None = None
if name is not None:
rename_record(session, reference_id, name)
if user_metadata is not None:
new_meta = dict(user_metadata)
elif computed_filename:
current_meta = ref.user_metadata or {}
if current_meta.get("filename") != computed_filename:
new_meta = dict(current_meta)
if new_meta is not None:
if computed_filename:
new_meta["filename"] = computed_filename
set_reference_metadata(
session, reference_id=reference_id, user_metadata=new_meta
session.execute(
update(Asset)
.where(Asset.id == reference_id)
.values(user_metadata=dict(user_metadata), updated_at=get_utc_now())
)
touched = True
manual_tags_before: set[str] = set()
if tags is not None:
set_reference_tags(
session,
reference_id=reference_id,
tags=tags,
origin=tag_origin,
manual_tags_before = _fetch_manual_tags(session, reference_id)
session.execute(
delete(AssetTag).where(
AssetTag.asset_id == reference_id,
AssetTag.origin != "automatic",
)
)
touched = True
if mime_type is not None:
updated = update_asset_hash_and_mime(
session, asset_id=ref.asset_id, mime_type=mime_type
session.execute(
update(Asset)
.where(Asset.id == reference_id)
.values(mime_type=mime_type, updated_at=get_utc_now())
)
if updated:
touched = True
if preview_id is not None:
set_reference_preview(
session,
reference_id=reference_id,
preview_reference_id=preview_id,
if session.get(Asset, preview_id) is None:
raise ValueError(
f"preview_id {preview_id!r} does not reference an existing asset"
)
session.execute(
update(Asset)
.where(Asset.id == reference_id)
.values(preview_id=preview_id, updated_at=get_utc_now())
)
touched = True
if touched and user_metadata is None:
update_reference_updated_at(session, reference_id=reference_id)
result = fetch_reference_asset_and_tags(
session,
reference_id=reference_id,
owner_id=owner_id,
)
if not result:
raise RuntimeError("State changed during update")
ref, asset, tag_list = result
detail = AssetDetailResult(
ref=extract_reference_data(ref),
asset=extract_asset_data(asset),
tags=tag_list,
)
if tags is not None:
for tag_name in normalize_tags(list(tags)):
ensure_tag(session, tag_name)
ensure_tag_link(
session,
asset_id=reference_id,
tag_name=tag_name,
origin=tag_origin,
)
session.flush()
if _fetch_manual_tags(session, reference_id) != manual_tags_before:
bump_record_updated_at(session, reference_id)
session.commit()
return detail
detail = get_asset_detail(reference_id)
if detail is None:
raise RuntimeError("Asset deleted during update")
return detail
def delete_asset_reference(
reference_id: str,
owner_id: str,
delete_content_if_orphan: bool = True,
) -> bool:
"""Delete an asset reference.
With ``delete_content_if_orphan=False`` (a soft delete), the reference is
hidden and the underlying content is preserved. With ``True``, the content
is also removed once it becomes orphaned.
Note: the public DELETE /api/assets/{id} endpoint always soft-deletes
(passes ``False``); the orphan-reclamation path is intentionally
internal-only, retained for a future GC/admin caller.
"""
with create_session() as session:
if not delete_content_if_orphan:
# Soft delete: mark the reference as deleted but keep everything
deleted = soft_delete_reference_by_id(
session, reference_id=reference_id, owner_id=owner_id
)
session.commit()
return deleted
ref_row = get_reference_by_id(session, reference_id=reference_id)
asset_id = ref_row.asset_id if ref_row else None
file_path = ref_row.file_path if ref_row else None
deleted = delete_reference_by_id(
session, reference_id=reference_id, owner_id=owner_id
)
if not deleted:
session.commit()
if get_record_by_id(session, reference_id) is None:
return False
if not asset_id:
session.commit()
return True
still_exists = reference_exists_for_asset_id(session, asset_id=asset_id)
if still_exists:
session.commit()
return True
# Orphaned asset - gather ALL file paths (including
# soft-deleted / missing refs) so their on-disk files get cleaned up.
file_paths = list_all_file_paths_by_asset_id(session, asset_id=asset_id)
# Also include the just-deleted file path
if file_path:
file_paths.append(file_path)
asset_row = session.get(Asset, asset_id)
if asset_row is not None:
session.delete(asset_row)
delete_record(session, reference_id)
session.commit()
# Delete files after commit
for p in file_paths:
with contextlib.suppress(Exception):
if p and os.path.isfile(p):
os.remove(p)
return True
def set_asset_preview(
reference_id: str,
preview_reference_id: str | None = None,
owner_id: str = "",
) -> AssetDetailResult:
with create_session() as session:
get_reference_with_owner_check(session, reference_id, owner_id)
set_reference_preview(
session,
reference_id=reference_id,
preview_reference_id=preview_reference_id,
)
result = fetch_reference_asset_and_tags(
session, reference_id=reference_id, owner_id=owner_id
)
if not result:
raise RuntimeError("State changed during preview update")
ref, asset, tags = result
detail = AssetDetailResult(
ref=extract_reference_data(ref),
asset=extract_asset_data(asset),
tags=tags,
)
session.commit()
return detail
return True
def asset_exists(asset_hash: str) -> bool:
try:
canonical = validate_blake3_hash(asset_hash)
except ValueError:
return False
with create_session() as session:
return asset_exists_by_hash(session, asset_hash=asset_hash)
def get_asset_by_hash(asset_hash: str) -> AssetData | None:
with create_session() as session:
asset = queries_get_asset_by_hash(session, asset_hash=asset_hash)
return extract_asset_data(asset)
# Sort fields that support cursor pagination. `last_access_time` is not
# in this list — it falls back to offset/limit.
_CURSOR_SORT_FIELDS = ("created_at", "updated_at", "name", "size")
def list_assets_page(
owner_id: str = "",
include_tags: Sequence[str] | None = None,
exclude_tags: Sequence[str] | None = None,
name_contains: str | None = None,
metadata_filter: dict | None = None,
limit: int = 20,
offset: int = 0,
sort: str = "created_at",
order: str = "desc",
after: str | None = None,
# Appended last so pre-existing positional callers keep binding correctly.
any_tags: Sequence[str] | None = None,
) -> ListAssetsResult:
"""List assets with optional cursor pagination.
When ``after`` is supplied it overrides ``offset``. The cursor's sort field
must match ``sort`` and be in the cursor-supported allowlist; mismatches
raise InvalidCursorError so the handler can map to 400 INVALID_CURSOR.
"""
cursor_value: object | None = None
cursor_id: str | None = None
# Mint next_cursor on every page where the sort is cursor-supported, not
# only when the request itself arrived with a cursor. Otherwise a first
# request (no `after`) returns next_cursor=None and the client can never
# enter cursor mode.
mint_cursor = sort in _CURSOR_SORT_FIELDS
if after is not None:
if sort not in _CURSOR_SORT_FIELDS:
raise InvalidCursorError(
f"cursor pagination is not supported for sort={sort!r}"
)
payload = decode_cursor(after, _CURSOR_SORT_FIELDS, expected_order=order)
if payload.sort_field != sort:
raise InvalidCursorError(
f"cursor sort field {payload.sort_field!r} does not match request sort {sort!r}"
)
cursor_value, cursor_id = _resolve_cursor_value(payload), payload.id
# Over-fetch by one row so we can distinguish "exactly `limit` rows total
# remaining" from "more rows past this page" without a second query. Drop
# the sentinel before returning.
fetch_limit = limit + 1 if mint_cursor else limit
with create_session() as session:
refs, tag_map, total = list_references_page(
session,
owner_id=owner_id,
include_tags=include_tags,
exclude_tags=exclude_tags,
any_tags=any_tags,
name_contains=name_contains,
metadata_filter=metadata_filter,
limit=fetch_limit,
offset=offset,
sort=sort,
order=order,
after_cursor_value=cursor_value,
after_cursor_id=cursor_id,
)
next_cursor: str | None = None
if mint_cursor and len(refs) > limit:
# There's at least one more row past this page — mint a cursor from
# the last row of the page (i.e. index `limit - 1`, since we
# over-fetched), and drop the sentinel.
next_cursor = _encode_next_cursor(refs[limit - 1], sort, order)
refs = refs[:limit]
items: list[AssetSummaryData] = []
for ref in refs:
items.append(
AssetSummaryData(
ref=extract_reference_data(ref),
asset=extract_asset_data(ref.asset),
tags=tag_map.get(ref.id, []),
)
)
return ListAssetsResult(items=items, total=total, next_cursor=next_cursor)
def _resolve_cursor_value(payload: CursorPayload) -> object:
"""Map a decoded cursor payload to a column-typed Python value."""
if payload.sort_field in ("created_at", "updated_at"):
# DB stores naive UTC; strip tzinfo so the comparison binds against a
# `TIMESTAMP WITHOUT TIME ZONE` column without an offset shift.
return decode_cursor_time(payload).replace(tzinfo=None)
if payload.sort_field == "size":
return decode_cursor_int(payload)
return payload.value # name, str-typed
def _encode_next_cursor(ref, sort: str, order: str) -> str | None:
"""Mint a cursor pointing at *ref* for the given sort dimension.
Returns None when the boundary row carries a NULL sort value (e.g. an asset
record whose size_bytes hasn't been backfilled). Continuing pagination
across a NULL boundary is undefined under keyset ordering — better to
truncate cleanly here than to mint a cursor that mis-positions.
"""
if sort == "name":
return encode_cursor("name", ref.name, ref.id, order=order)
if sort == "size":
if ref.asset is None or ref.asset.size_bytes is None:
return None
return encode_cursor("size", str(ref.asset.size_bytes), ref.id, order=order)
# created_at / updated_at — DB datetimes are naive UTC; attach tz before encoding.
value = ref.created_at if sort == "created_at" else ref.updated_at
if value is None:
return None
return encode_cursor_from_time(sort, value.replace(tzinfo=timezone.utc), ref.id, order=order)
return lookup_for_view(session, canonical) is not None
def resolve_hash_to_path(
asset_hash: str,
owner_id: str = "",
) -> DownloadResolutionResult | None:
"""Resolve a blake3 hash to an on-disk file path.
"""Resolve a blake3 hash to an on-disk file path via lookup_for_view.
Only references visible to *owner_id* are considered (owner-less
references are always visible).
Uses the first qualified live content row. Temp paths are excluded from all
hash lookups inside qualified_content_iterator, so a hash resolving only to
temp content returns None. Updates last_access_time on every record pointing
at the served content.
Returns a DownloadResolutionResult with abs_path, content_type, and
download_name, or None if no asset or live path is found.
Filename and Content-Type both come from the newest record so they never
describe different records. Deleting the last record preserves its content,
so content with zero records stays servable off the content path alone.
"""
try:
canonical = validate_blake3_hash(asset_hash)
except ValueError:
return None
with create_session() as session:
asset = queries_get_asset_by_hash(session, asset_hash)
if not asset:
content = lookup_for_view(session, canonical)
if content is None:
return None
refs = list_references_by_asset_id(session, asset_id=asset.id)
visible = [
r for r in refs
if r.owner_id == "" or r.owner_id == owner_id
]
abs_path = select_best_live_path(visible)
if not abs_path:
return None
display_name = os.path.basename(abs_path)
for ref in visible:
if ref.file_path == abs_path and ref.name:
display_name = ref.name
break
records = list(
session.scalars(
select(Asset)
.where(Asset.content_id == content.id)
.order_by(Asset.created_at, Asset.id)
)
)
display_name = os.path.basename(content.path)
mime_type = None
if records:
latest_record = records[-1]
display_name = latest_record.name or display_name
mime_type = latest_record.mime_type
for record in records:
update_record_access_time(session, record.id)
abs_path = content.path
session.commit()
ctype = (
asset.mime_type
mime_type
or mimetypes.guess_type(display_name)[0]
or mimetypes.guess_type(abs_path)[0]
or "application/octet-stream"
)
return DownloadResolutionResult(
@@ -430,40 +230,33 @@ def get_preview_file_paths(preview_ids: list[str]) -> dict[str, str]:
if not preview_ids:
return {}
with create_session() as session:
return get_reference_paths_by_ids(session, reference_ids=preview_ids)
return get_preview_file_paths_by_ids(session, preview_ids=preview_ids)
def resolve_asset_for_download(
reference_id: str,
owner_id: str = "",
) -> DownloadResolutionResult:
with create_session() as session:
pair = fetch_reference_and_asset(
session, reference_id=reference_id, owner_id=owner_id
)
if not pair:
record = get_record_by_id(session, reference_id)
if record is None:
raise ValueError(f"AssetReference {reference_id} not found")
ref, asset = pair
content = session.get(AssetContent, record.content_id)
if (
content is None
or content.is_missing
or not os.path.isfile(content.path)
):
raise FileNotFoundError(
f"No live content for AssetReference {reference_id} "
f"(content id={record.content_id}, name={record.name})"
)
# For references with file_path, use that directly
if ref.file_path and os.path.isfile(ref.file_path):
abs_path = ref.file_path
else:
# For API-created refs without file_path, find a path from other refs
refs = list_references_by_asset_id(session, asset_id=asset.id)
abs_path = select_best_live_path(refs)
if not abs_path:
raise FileNotFoundError(
f"No live path for AssetReference {reference_id} "
f"(asset id={asset.id}, name={ref.name})"
)
ref_name = record.name
asset_mime = record.mime_type
abs_path = content.path
# Capture ORM attributes before commit (commit expires loaded objects)
ref_name = ref.name
asset_mime = asset.mime_type
update_reference_access_time(session, reference_id=reference_id)
update_record_access_time(session, reference_id)
session.commit()
ctype = (

View File

@@ -1,294 +0,0 @@
from __future__ import annotations
import os
import uuid
from dataclasses import dataclass
from datetime import datetime
from typing import TYPE_CHECKING, Any, TypedDict
from sqlalchemy.orm import Session
from app.assets.database.queries import (
bulk_insert_assets,
bulk_insert_references_ignore_conflicts,
bulk_insert_tags_and_meta,
delete_assets_by_ids,
get_existing_asset_ids,
get_reference_ids_by_ids,
get_references_by_paths_and_asset_ids,
get_unreferenced_unhashed_asset_ids,
restore_references_by_paths,
)
from app.assets.helpers import get_utc_now
if TYPE_CHECKING:
from app.assets.services.metadata_extract import ExtractedMetadata
class SeedAssetSpec(TypedDict):
"""Spec for seeding an asset from filesystem."""
abs_path: str
size_bytes: int
mtime_ns: int
info_name: str
tags: list[str]
fname: str
metadata: ExtractedMetadata | None
hash: str | None
mime_type: str | None
job_id: str | None
class AssetRow(TypedDict):
"""Row data for inserting an Asset."""
id: str
hash: str | None
size_bytes: int
mime_type: str | None
created_at: datetime
class ReferenceRow(TypedDict):
"""Row data for inserting an AssetReference."""
id: str
asset_id: str
file_path: str
loader_path: str | None
mtime_ns: int
owner_id: str
name: str
preview_id: str | None
user_metadata: dict[str, Any] | None
job_id: str | None
created_at: datetime
updated_at: datetime
last_access_time: datetime
class TagRow(TypedDict):
"""Row data for inserting a Tag."""
asset_reference_id: str
tag_name: str
origin: str
added_at: datetime
class MetadataRow(TypedDict):
"""Row data for inserting asset metadata."""
asset_reference_id: str
key: str
ordinal: int
val_str: str | None
val_num: float | None
val_bool: bool | None
val_json: dict[str, Any] | None
@dataclass
class BulkInsertResult:
"""Result of bulk asset insertion."""
inserted_refs: int
won_paths: int
lost_paths: int
def batch_insert_seed_assets(
session: Session,
specs: list[SeedAssetSpec],
owner_id: str = "",
) -> BulkInsertResult:
"""Seed assets from filesystem specs in batch.
Each spec is a dict with keys:
- abs_path: str
- size_bytes: int
- mtime_ns: int
- info_name: str
- tags: list[str]
- fname: Optional[str]
This function orchestrates:
1. Insert seed Assets (hash=NULL)
2. Claim references with ON CONFLICT DO NOTHING on file_path
3. Query to find winners (paths where our asset_id was inserted)
4. Delete Assets for losers (path already claimed by another asset)
5. Insert tags and metadata for successfully inserted references
Returns:
BulkInsertResult with inserted_refs, won_paths, lost_paths
"""
if not specs:
return BulkInsertResult(inserted_refs=0, won_paths=0, lost_paths=0)
current_time = get_utc_now()
asset_rows: list[AssetRow] = []
reference_rows: list[ReferenceRow] = []
path_to_asset_id: dict[str, str] = {}
asset_id_to_ref_data: dict[str, dict] = {}
absolute_path_list: list[str] = []
for spec in specs:
absolute_path = os.path.abspath(spec["abs_path"])
existing_asset_id = path_to_asset_id.get(absolute_path)
if existing_asset_id is not None:
existing_tags = asset_id_to_ref_data[existing_asset_id]["tags"]
asset_id_to_ref_data[existing_asset_id]["tags"] = list(
dict.fromkeys([*existing_tags, *spec["tags"]])
)
continue
asset_id = str(uuid.uuid4())
reference_id = str(uuid.uuid4())
absolute_path_list.append(absolute_path)
path_to_asset_id[absolute_path] = asset_id
mime_type = spec.get("mime_type")
asset_rows.append(
{
"id": asset_id,
"hash": spec.get("hash"),
"size_bytes": spec["size_bytes"],
"mime_type": mime_type,
"created_at": current_time,
}
)
# Build user_metadata from extracted metadata or fallback to filename
extracted_metadata = spec.get("metadata")
if extracted_metadata:
user_metadata: dict[str, Any] | None = extracted_metadata.to_user_metadata()
elif spec["fname"]:
user_metadata = {"filename": spec["fname"]}
else:
user_metadata = None
reference_rows.append(
{
"id": reference_id,
"asset_id": asset_id,
"file_path": absolute_path,
# spec["fname"] is compute_loader_path(abs_path) from build_asset_specs.
"loader_path": spec["fname"],
"mtime_ns": spec["mtime_ns"],
"owner_id": owner_id,
"name": spec["info_name"],
"preview_id": None,
"user_metadata": user_metadata,
"job_id": spec.get("job_id"),
"created_at": current_time,
"updated_at": current_time,
"last_access_time": current_time,
}
)
asset_id_to_ref_data[asset_id] = {
"reference_id": reference_id,
"tags": spec["tags"],
"filename": spec["fname"],
"extracted_metadata": extracted_metadata,
}
bulk_insert_assets(session, asset_rows)
# Filter reference rows to only those whose assets were actually inserted
# (assets with duplicate hashes are silently dropped by ON CONFLICT DO NOTHING)
inserted_asset_ids = get_existing_asset_ids(
session, [r["asset_id"] for r in reference_rows]
)
reference_rows = [r for r in reference_rows if r["asset_id"] in inserted_asset_ids]
bulk_insert_references_ignore_conflicts(session, reference_rows)
restore_references_by_paths(session, absolute_path_list)
winning_paths = get_references_by_paths_and_asset_ids(session, path_to_asset_id)
inserted_paths = {
path
for path in absolute_path_list
if path_to_asset_id[path] in inserted_asset_ids
}
losing_paths = inserted_paths - winning_paths
lost_asset_ids = [path_to_asset_id[path] for path in losing_paths]
if lost_asset_ids:
delete_assets_by_ids(session, lost_asset_ids)
if not winning_paths:
return BulkInsertResult(
inserted_refs=0,
won_paths=0,
lost_paths=len(losing_paths),
)
# Get reference IDs for winners
winning_ref_ids = [
asset_id_to_ref_data[path_to_asset_id[path]]["reference_id"]
for path in winning_paths
]
inserted_ref_ids = get_reference_ids_by_ids(session, winning_ref_ids)
tag_rows: list[TagRow] = []
metadata_rows: list[MetadataRow] = []
if inserted_ref_ids:
for path in winning_paths:
asset_id = path_to_asset_id[path]
ref_data = asset_id_to_ref_data[asset_id]
ref_id = ref_data["reference_id"]
if ref_id not in inserted_ref_ids:
continue
for tag in ref_data["tags"]:
tag_rows.append(
{
"asset_reference_id": ref_id,
"tag_name": tag,
"origin": "automatic",
"added_at": current_time,
}
)
# Use extracted metadata for meta rows if available
extracted_metadata = ref_data.get("extracted_metadata")
if extracted_metadata:
metadata_rows.extend(extracted_metadata.to_meta_rows(ref_id))
elif ref_data["filename"]:
# Fallback: just store filename
metadata_rows.append(
{
"asset_reference_id": ref_id,
"key": "filename",
"ordinal": 0,
"val_str": ref_data["filename"],
"val_num": None,
"val_bool": None,
"val_json": None,
}
)
bulk_insert_tags_and_meta(session, tag_rows=tag_rows, meta_rows=metadata_rows)
return BulkInsertResult(
inserted_refs=len(inserted_ref_ids),
won_paths=len(winning_paths),
lost_paths=len(losing_paths),
)
def cleanup_unreferenced_assets(session: Session) -> int:
"""Hard-delete unhashed assets with no active references.
This is a destructive operation intended for explicit cleanup.
Only deletes assets where hash=None and all references are missing.
Returns:
Number of assets deleted
"""
unreferenced_ids = get_unreferenced_unhashed_asset_ids(session)
return delete_assets_by_ids(session, unreferenced_ids)

View File

@@ -0,0 +1,189 @@
"""Tracks the persisted hashing mode and carries the existing catalog across an
off-to-on switch. Turning hashing on enqueues every live content row for
verification and the queue is drained in the background; a path that cannot be
read is retried a bounded number of times and then has its stored hash cleared,
so the queue can empty and the mode can flip. An unverifiable digest never
survives into the persisted enabled state.
"""
from __future__ import annotations
import logging
import os
from collections import deque
from dataclasses import dataclass
from typing import Final
from sqlalchemy import select
from sqlalchemy.orm import Session
from app.assets import mode as _mode
from app.assets.database.models import AssetContent, AssetSystemState
from app.assets.database.queries.records import create_content, create_record, mark_content_missing
from app.assets.helpers import to_stored_hash
from app.assets.services.path_utils import compute_loader_path, get_name_and_tags_from_asset_path
from app.assets.services.snapshot_hash import snapshot_hash
_KEY = "hash_mode"
_MAX_VERIFY_ATTEMPTS: Final = 3
@dataclass
class _PendingEntry:
path: str
ticks: int = 0
_PENDING_QUEUE: deque[_PendingEntry] = deque()
_PENDING_PATHS: set[str] = set()
_off_to_on_transition_in_flight = False
def clear_transition_queue() -> None:
global _off_to_on_transition_in_flight
_PENDING_QUEUE.clear()
_PENDING_PATHS.clear()
_off_to_on_transition_in_flight = False
def pending_transition_count() -> int:
return len(_PENDING_QUEUE)
def read_stored_mode(session: Session) -> str | None:
row = session.get(AssetSystemState, _KEY)
return row.value if row else None
def write_stored_mode(session: Session, value: str) -> None:
row = session.get(AssetSystemState, _KEY)
if row is None:
session.add(AssetSystemState(key=_KEY, value=value))
else:
row.value = value
session.flush()
def record_transition_intent(session: Session) -> str | None:
stored = read_stored_mode(session)
runtime = "on" if _mode.hashing_enabled() else "off"
if stored is None:
write_stored_mode(session, runtime)
return None
if stored == "off" and runtime == "on":
return "off_to_on"
if stored == "on" and runtime == "off":
write_stored_mode(session, "off")
return "on_to_off"
return None
def enqueue_transition_work(session: Session, transition: str | None) -> None:
global _off_to_on_transition_in_flight
if transition != "off_to_on":
return
_off_to_on_transition_in_flight = True
rows = session.scalars(
select(AssetContent).where(AssetContent.is_missing.is_(False))
)
for row in rows:
if row.path not in _PENDING_PATHS:
_PENDING_QUEUE.append(_PendingEntry(row.path))
_PENDING_PATHS.add(row.path)
def _retry_or_retire(session: Session, entry: _PendingEntry) -> None:
entry.ticks += 1
if entry.ticks < _MAX_VERIFY_ATTEMPTS:
_PENDING_QUEUE.append(entry)
_PENDING_PATHS.add(entry.path)
return
content = session.scalars(
select(AssetContent).where(
AssetContent.path == entry.path, AssetContent.is_missing.is_(False)
)
).first()
if content is not None:
content.hash = None
logging.warning(
"Could not verify %s in %d attempts; clearing its stored hash so the hash-mode "
"transition can complete",
entry.path,
_MAX_VERIFY_ATTEMPTS,
)
def drain_transition_queue(session: Session) -> None:
global _off_to_on_transition_in_flight
pending_count = len(_PENDING_QUEUE)
for _ in range(pending_count):
entry = _PENDING_QUEUE.popleft()
_PENDING_PATHS.discard(entry.path)
path = entry.path
try:
snapshot = snapshot_hash(path)
except OSError:
_retry_or_retire(session, entry)
continue
if snapshot is None:
# snapshot_hash returns None for vanished and unstable files; stat distinguishes them.
try:
os.stat(path)
except FileNotFoundError:
gone = session.scalars(
select(AssetContent).where(
AssetContent.path == path, AssetContent.is_missing.is_(False)
)
).first()
if gone is not None:
mark_content_missing(session, gone.id)
except OSError:
_retry_or_retire(session, entry)
else:
_retry_or_retire(session, entry)
continue
digest, stat = snapshot
stored_hash = to_stored_hash(digest)
content = session.scalars(
select(AssetContent).where(
AssetContent.path == path, AssetContent.is_missing.is_(False)
)
).first()
if content is None:
continue
if content.hash is None:
content.hash = stored_hash
content.size_bytes = stat.st_size
content.mtime_ns = stat.st_mtime_ns
elif content.hash != stored_hash:
try:
name, tags = get_name_and_tags_from_asset_path(path)
except ValueError:
logging.warning(
"Skipping hash-mode split for out-of-root path: %s", path
)
continue
mark_content_missing(session, content.id)
replacement = create_content(
session,
path=path,
hash=stored_hash,
size_bytes=stat.st_size,
mtime_ns=stat.st_mtime_ns,
)
create_record(
session,
content_id=replacement.id,
name=name,
loader_path=compute_loader_path(path),
tags=tags,
)
else:
content.size_bytes = stat.st_size
content.mtime_ns = stat.st_mtime_ns
if _off_to_on_transition_in_flight and not _PENDING_QUEUE:
write_stored_mode(session, "on")
_off_to_on_transition_in_flight = False

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,96 @@
"""Resolves a hash to a content row that can actually be served, and claims that
row so a concurrent writer cannot retire it mid-use. A row qualifies only if
its file exists, sits outside the temp directory, and its recorded size and
mtime still match the file — an unrecorded mtime excuses the mtime comparison
but never the size comparison.
"""
from __future__ import annotations
import os
from collections.abc import Iterator
from pathlib import Path
import folder_paths
from sqlalchemy import select, update
from sqlalchemy.orm import Session
from app.assets import mode
from app.assets.database.models import AssetContent
def is_temp_path(path: str) -> bool:
try:
temp_root = Path(os.path.abspath(folder_paths.get_temp_directory()))
candidate = Path(os.path.abspath(path))
return candidate.is_relative_to(temp_root)
except OSError:
return False
def _stat_consistent(content: AssetContent) -> bool:
try:
stat = os.stat(content.path)
except FileNotFoundError:
return False
if content.mtime_ns is not None and stat.st_mtime_ns != content.mtime_ns:
return False
if stat.st_size != content.size_bytes:
return False
return True
def _qualifies(content: AssetContent) -> bool:
return (
not content.is_missing
and os.path.isfile(content.path)
and _stat_consistent(content)
and not is_temp_path(content.path)
)
def qualified_content_iterator(session: Session, hash: str) -> Iterator[AssetContent]:
rows = session.scalars(
select(AssetContent)
.where(AssetContent.hash == hash, AssetContent.is_missing.is_(False))
.order_by(AssetContent.created_at, AssetContent.id)
)
for row in rows:
if _qualifies(row):
yield row
def claim_qualified_content(session: Session, content_id: str, hash: str) -> bool:
"""Claim a live matching content row before attaching a record.
The conditional update takes this session's SQLite write lock through commit.
False means the row was retired or changed after lookup.
"""
# This session's connection keeps the claim's write lock through this session's commit.
result = session.connection().execute(
update(AssetContent)
.where(
AssetContent.id == content_id,
AssetContent.hash == hash,
AssetContent.is_missing.is_(False),
)
.values(is_missing=False)
)
return result.rowcount == 1
def refresh_qualified_content(session: Session, content_id: str) -> AssetContent | None:
content = session.get(AssetContent, content_id, populate_existing=True)
if content is None or not _qualifies(content):
return None
return content
def lookup_for_from_hash(session: Session, hash: str) -> AssetContent | None:
if not mode.hashing_enabled():
return None
return next(qualified_content_iterator(session, hash), None)
def lookup_for_view(session: Session, hash: str) -> AssetContent | None:
return next(qualified_content_iterator(session, hash), None)

View File

@@ -1,9 +1,13 @@
"""Defines the frozen result types the service layer hands back to routes and to
the seeder, keeping those callers off ORM objects and the sessions that own
them. Each type carries only what its caller renders, so nothing can be lazily
loaded from a row after its session has closed.
"""
from dataclasses import dataclass
from datetime import datetime
from typing import Any, NamedTuple
from app.assets.database.models import Asset, AssetReference
UserMetadata = dict[str, Any] | None
@@ -12,6 +16,7 @@ class AssetData:
hash: str | None
size_bytes: int | None
mime_type: str | None
is_missing: bool = False
@dataclass(frozen=True)
@@ -81,35 +86,27 @@ class DownloadResolutionResult:
download_name: str
@dataclass(frozen=True)
class RegisteredAsset:
id: str
content_id: str
job_id: str | None
name: str
@dataclass(frozen=True)
class UploadAssetView:
asset: RegisteredAsset
asset_hash: str | None
size: int | None
mime_type: str | None
tags: list[str]
@dataclass(frozen=True)
class UploadResult:
ref: ReferenceData
content_id: str
asset: AssetData
tags: list[str]
created_new: bool
def extract_reference_data(ref: AssetReference) -> ReferenceData:
return ReferenceData(
id=ref.id,
name=ref.name,
file_path=ref.file_path,
loader_path=ref.loader_path,
user_metadata=ref.user_metadata,
preview_id=ref.preview_id,
system_metadata=ref.system_metadata,
job_id=ref.job_id,
created_at=ref.created_at,
updated_at=ref.updated_at,
last_access_time=ref.last_access_time,
)
def extract_asset_data(asset: Asset | None) -> AssetData | None:
if asset is None:
return None
return AssetData(
hash=asset.hash,
size_bytes=asset.size_bytes,
mime_type=asset.mime_type,
)

View File

@@ -0,0 +1,51 @@
"""Hashes a file and returns the stat it proved describes those exact bytes.
Identity, size and mtime are sampled before the read, on the open handle at
both ends of it, and once more afterwards; if any sample disagrees the file
moved under the reader and the result is discarded rather than returned.
Callers persist the stat that comes back, which is what makes a stored hash and
a stored size describe one observation.
"""
from __future__ import annotations
import os
from dataclasses import dataclass
from blake3 import blake3
@dataclass(frozen=True, slots=True)
class _Snapshot:
dev: int
ino: int
mtime_ns: int
size: int
def _snapshot(stat_result: os.stat_result) -> _Snapshot:
return _Snapshot(
dev=stat_result.st_dev,
ino=stat_result.st_ino,
mtime_ns=stat_result.st_mtime_ns,
size=stat_result.st_size,
)
def snapshot_hash(
path: str, chunk_size: int = 8 * 1024 * 1024
) -> tuple[str, os.stat_result] | None:
try:
pre_stat = _snapshot(os.stat(path))
hasher = blake3()
with open(path, "rb") as file:
open_stat = _snapshot(os.fstat(file.fileno()))
while chunk := file.read(chunk_size):
hasher.update(chunk)
post_hash_stat = _snapshot(os.fstat(file.fileno()))
post_stat_result = os.stat(path)
except FileNotFoundError:
return None
post_stat = _snapshot(post_stat_result)
if len({pre_stat, open_stat, post_hash_stat, post_stat}) != 1:
return None
return hasher.hexdigest(), post_stat_result

View File

@@ -1,14 +1,27 @@
"""Applies and removes tags on a record and answers tag listing and histogram
queries. Adding reports what it actually inserted, read back after the write,
so a caller that lost a race is told the tag was already present rather than
that it added one. Removal leaves automatically applied tags in place and
reports them as protected instead of silently skipping them.
"""
from typing import Sequence
from sqlalchemy import delete, select
from app.assets.database.queries import (
AddTagsResult,
RemoveTagsResult,
add_tags_to_reference,
get_reference_with_owner_check,
list_tags_with_usage,
remove_tags_from_reference,
)
from app.assets.database.queries.records import (
bump_record_updated_at,
ensure_tag,
ensure_tag_link,
)
from app.assets.database.queries.tags import list_tag_counts_for_filtered_assets
from app.assets.database.models import Asset, AssetTag
from app.assets.helpers import normalize_tags
from app.assets.services.schemas import TagUsage
from app.database.db import create_session
@@ -17,40 +30,93 @@ def apply_tags(
reference_id: str,
tags: list[str],
origin: str = "manual",
owner_id: str = "",
) -> AddTagsResult:
with create_session() as session:
ref_row = get_reference_with_owner_check(session, reference_id, owner_id)
if session.get(Asset, reference_id) is None:
raise ValueError(f"Asset {reference_id} not found")
result = add_tags_to_reference(
session,
reference_id=reference_id,
tags=tags,
origin=origin,
create_if_missing=True,
reference_row=ref_row,
normalized_tags = normalize_tags(tags)
requested_tags = set(normalized_tags)
added: list[str] = []
for tag_name in normalized_tags:
ensure_tag(session, tag_name)
if ensure_tag_link(
session,
asset_id=reference_id,
tag_name=tag_name,
origin=origin,
):
added.append(tag_name)
session.flush()
if added:
bump_record_updated_at(session, reference_id)
total_tags = list(
session.scalars(
select(AssetTag.tag_name)
.where(AssetTag.asset_id == reference_id)
.order_by(AssetTag.tag_name)
)
)
session.commit()
return result
return AddTagsResult(
added=sorted(added),
already_present=sorted((requested_tags & set(total_tags)) - set(added)),
total_tags=total_tags,
)
def remove_tags(
reference_id: str,
tags: list[str],
owner_id: str = "",
) -> RemoveTagsResult:
with create_session() as session:
get_reference_with_owner_check(session, reference_id, owner_id)
if session.get(Asset, reference_id) is None:
raise ValueError(f"Asset {reference_id} not found")
result = remove_tags_from_reference(
session,
reference_id=reference_id,
tags=tags,
requested_tags = set(normalize_tags(tags))
removable_tags = set(
session.scalars(
select(AssetTag.tag_name).where(
AssetTag.asset_id == reference_id,
AssetTag.origin != "automatic",
AssetTag.tag_name.in_(requested_tags),
)
)
)
protected_tags = set(
session.scalars(
select(AssetTag.tag_name).where(
AssetTag.asset_id == reference_id,
AssetTag.origin == "automatic",
AssetTag.tag_name.in_(requested_tags),
)
)
)
if removable_tags:
session.execute(
delete(AssetTag).where(
AssetTag.asset_id == reference_id,
AssetTag.origin != "automatic",
AssetTag.tag_name.in_(removable_tags),
)
)
bump_record_updated_at(session, reference_id)
total_tags = list(
session.scalars(
select(AssetTag.tag_name)
.where(AssetTag.asset_id == reference_id)
.order_by(AssetTag.tag_name)
)
)
session.commit()
return result
return RemoveTagsResult(
removed=sorted(removable_tags),
not_present=sorted(requested_tags - removable_tags - protected_tags),
total_tags=total_tags,
protected=sorted(protected_tags),
)
def list_tags(
@@ -59,7 +125,6 @@ def list_tags(
offset: int = 0,
order: str = "count_desc",
include_zero: bool = True,
owner_id: str = "",
) -> tuple[list[TagUsage], int]:
limit = max(1, min(1000, limit))
offset = max(0, offset)
@@ -72,18 +137,15 @@ def list_tags(
offset=offset,
include_zero=include_zero,
order=order,
owner_id=owner_id,
)
return [TagUsage(name, count) for name, count in rows], total
def list_tag_histogram(
owner_id: str = "",
include_tags: Sequence[str] | None = None,
exclude_tags: Sequence[str] | None = None,
name_contains: str | None = None,
metadata_filter: dict | None = None,
limit: int = 100,
# Appended last so pre-existing positional callers keep binding correctly.
any_tags: Sequence[str] | None = None,
@@ -91,11 +153,9 @@ def list_tag_histogram(
with create_session() as session:
return list_tag_counts_for_filtered_assets(
session,
owner_id=owner_id,
include_tags=include_tags,
exclude_tags=exclude_tags,
any_tags=any_tags,
name_contains=name_contains,
metadata_filter=metadata_filter,
limit=limit,
)

View File

@@ -21,6 +21,7 @@ try:
from app.database.models import Base
import app.assets.database.models # noqa: F401 — register models with Base.metadata
import blake3 # noqa: F401 — verify the hard dependency is importable at startup
_DB_AVAILABLE = True
except ImportError as e:
@@ -184,6 +185,29 @@ def _init_file_db(db_url):
prepare_file_db_path(db_path)
db_exists = os.path.exists(db_path)
# Lock BEFORE any migration work — deliberately diverging from upstream master, whose
# "it would block Alembic" rationale is false (the lock guards a separate `<db>.lock`
# file). Only this order makes revision inspection, backup, upgrade and the failure-path
# restore mutually exclusive between processes.
_acquire_file_lock(db_path)
try:
_migrate_and_bind(db_url, db_path, db_exists)
except Exception:
_db_lock.release()
raise
_DESTRUCTIVE_REVISION = "0007_record_content_split"
def _upgrade_discards_the_catalog(script, target_rev, current_rev):
return any(
revision.revision == _DESTRUCTIVE_REVISION
for revision in script.iterate_revisions(upper=target_rev, lower=current_rev)
)
def _migrate_and_bind(db_url, db_path, db_exists):
config = get_alembic_config()
# Check if we need to upgrade
@@ -225,11 +249,15 @@ def _init_file_db(db_url):
logging.exception("Error upgrading database: ")
raise e
# Acquire an OS-level file lock after migrations are complete.
# Alembic uses its own connection, so we must wait until it's done
# before locking — otherwise our own lock blocks the migration.
if backup_path and _upgrade_discards_the_catalog(script, target_rev, current_rev):
log_startup_warning(
f"The asset catalog was rebuilt from scratch by migration "
f"{_DESTRUCTIVE_REVISION}: manual tags, user metadata, previews, renames, "
f"API-created records and job_id links from the previous database were "
f"discarded. The database from before the upgrade was kept at {backup_path}."
)
conn.close()
_acquire_file_lock(db_path)
global Session
Session = sessionmaker(bind=engine)

View File

@@ -109,6 +109,7 @@ _CORE_FEATURE_FLAGS: dict[str, Any] = {
"max_upload_size": args.max_upload_size * 1024 * 1024, # Convert MB to bytes
"extension": {"manager": {"supports_v4": True}},
"node_replacements": True,
# Mirrors the constructed AssetManager; no degradation path exists, so this always agrees with it.
"assets": args.enable_assets,
}

View File

@@ -1,66 +1,115 @@
"""Enrich executed-node output entries with asset id."""
"""Registers the files a prompt produced as assets, at the moment the execution
result is emitted. Each output is matched to a path inside its declared base
directory and skipped when it escapes that directory or is not on disk, and a
run served from cache replays the same registration so cached results still
yield assets. Registering at emission rather than by inspecting the cache
afterwards keeps this independent of any cache eviction policy.
"""
import copy
import logging
import os
from typing import TYPE_CHECKING, Callable
if TYPE_CHECKING:
from app.assets.manager import AssetManager
from app.assets.services.schemas import RegisteredAsset
from comfy_execution.server_protocol import ExecutionServer
from execution import CacheEntry
def enrich_output_with_assets(output_ui: dict) -> dict:
"""Register file-type output entries as assets and inject their ``id``.
def _resolve_output_path(entry: dict) -> str | None:
"""Resolve an output entry to an absolute, in-base, on-disk file path.
Runs at output-processing time, once per produced output, when
--enable-assets is set. Returns a new dict; entries without a resolvable
on-disk file path are left unchanged. Errors are caught per-entry so a
failure never blocks execution or the other entries.
Returns ``None`` (skip, no registration) when the type is unknown, the
resolved path escapes its base directory, or the file does not exist.
"""
from comfy.cli_args import args
if not args.enable_assets:
return output_ui
import folder_paths
from app.assets.services.ingest import register_file_in_place, DependencyMissingError
enriched = {}
for key, entries in output_ui.items():
base = folder_paths.get_directory_by_type(entry["type"])
if base is None:
return None
base_abs = os.path.abspath(base)
abs_path = os.path.abspath(os.path.join(base_abs, entry.get("subfolder") or "", entry["filename"]))
try:
if os.path.commonpath([base_abs, abs_path]) != base_abs:
return None
except ValueError:
return None
if not os.path.isfile(abs_path):
return None
return abs_path
def _enrich_in_place(
output_ui: dict,
job_id: str | None,
register: Callable[[str, str | None], "RegisteredAsset | None"],
) -> None:
"""S10.6: producers that write the same output path are not coalesced (unsupported)."""
for entries in output_ui.values():
if not isinstance(entries, list):
enriched[key] = entries
continue
new_entries = []
for entry in entries:
if not isinstance(entry, dict) or "filename" not in entry or "type" not in entry:
new_entries.append(entry)
continue
try:
base = folder_paths.get_directory_by_type(entry["type"])
if base is None:
new_entries.append(entry)
abs_path = _resolve_output_path(entry)
if abs_path is None:
continue
base_abs = os.path.abspath(base)
abs_path = os.path.abspath(os.path.join(base_abs, entry.get("subfolder") or "", entry["filename"]))
try:
if os.path.commonpath([base_abs, abs_path]) != base_abs:
raise ValueError("escapes base")
except ValueError:
logging.warning("Asset enrichment skipped (path escapes base): %s", entry.get("filename"))
new_entries.append(entry)
continue
if not os.path.isfile(abs_path):
new_entries.append(entry)
continue
# Register unconditionally: the file was just produced, and
# register_file_in_place re-hashes so an overwritten path can
# never carry a stale id.
result = register_file_in_place(
abs_path=abs_path,
name=entry["filename"],
tags=[entry["type"]],
)
entry = dict(entry)
entry["id"] = result.ref.id
except DependencyMissingError:
logging.warning("Asset enrichment skipped (blake3 not available): %s", entry.get("filename"))
result = register(abs_path, job_id)
if result is not None:
entry["id"] = result.id
except Exception:
logging.warning("Failed to enrich output entry with asset id: %s", entry.get("filename"), exc_info=True)
new_entries.append(entry)
enriched[key] = new_entries
logging.warning("Asset registration failed for output: %s", entry.get("filename"), exc_info=True)
def _strip_ids(output_ui: dict) -> None:
for entries in output_ui.values():
if not isinstance(entries, list):
continue
for entry in entries:
if isinstance(entry, dict):
entry.pop("id", None)
def register_executed_outputs(output_ui: dict, job_id: str, asset_manager: "AssetManager") -> dict:
enriched = copy.deepcopy(output_ui)
if not asset_manager.enabled:
return enriched
_enrich_in_place(enriched, job_id, asset_manager.register_executed_output)
return enriched
def register_cached_outputs(ui_wrapper: dict | None, job_id: str, asset_manager: "AssetManager") -> dict | None:
if ui_wrapper is None:
return None
enriched = copy.deepcopy(ui_wrapper)
output_ui = enriched.get("output")
if not isinstance(output_ui, dict):
return enriched
_strip_ids(output_ui)
if not asset_manager.enabled:
return enriched
_enrich_in_place(output_ui, job_id, asset_manager.register_cached_output)
return enriched
def emit_cached_output(server: "ExecutionServer", node_id: str, display_node_id: str, cached: "CacheEntry", prompt_id: str, ui_outputs: dict, asset_manager: "AssetManager") -> None:
if node_id in ui_outputs:
return
enriched = register_cached_outputs(cached.ui, prompt_id, asset_manager)
if enriched is not None:
ui_outputs[node_id] = enriched
if server.client_id is None:
return
output = enriched.get("output") if enriched is not None else None
server.send_sync(
"executed",
{"node": node_id, "display_node": display_node_id, "output": output, "prompt_id": prompt_id},
server.client_id,
)

View File

@@ -7,6 +7,7 @@ from tqdm import tqdm
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from comfy_execution.graph import DynamicPrompt
from comfy_execution.server_protocol import ExecutionServer
from protocol import BinaryEventTypes
from comfy_api import feature_flags
@@ -150,7 +151,7 @@ class WebUIProgressHandler(ProgressHandler):
Handler that sends progress updates to the WebUI via WebSockets.
"""
def __init__(self, server_instance):
def __init__(self, server_instance: "ExecutionServer"):
super().__init__("webui")
self.server_instance = server_instance

View File

@@ -0,0 +1,11 @@
from typing import Protocol
class ExecutionServer(Protocol):
client_id: str | None
last_node_id: str | None
sockets_metadata: dict[str, dict[str, object]]
def send_sync(self, event: str | int, data: object, sid: str | None = None) -> None: ...
def queue_updated(self) -> None: ...

View File

@@ -7,7 +7,7 @@ import threading
import time
import traceback
from enum import Enum
from typing import List, Literal, NamedTuple, Optional, Union
from typing import TYPE_CHECKING, List, Literal, NamedTuple, Optional, Union
import asyncio
import torch
@@ -43,10 +43,14 @@ from comfy_execution.graph_utils import GraphBuilder, is_link
from comfy_execution.validation import validate_node_input
from comfy_execution.progress import get_progress_state, reset_progress_state, add_progress_handler, WebUIProgressHandler
from comfy_execution.utils import CurrentNodeContext
from comfy_execution.asset_enrichment import enrich_output_with_assets
from comfy_execution.asset_enrichment import register_executed_outputs, emit_cached_output
from comfy_api.internal import _ComfyNodeInternal, _NodeOutputInternal, first_real_override, is_class, make_locked_method_func
from comfy_api.latest import io, _io
from comfy_execution.cache_provider import _has_cache_providers, _get_cache_providers, _logger as _cache_logger
from app.assets.manager import AssetManager, default_asset_manager
if TYPE_CHECKING:
from comfy_execution.server_protocol import ExecutionServer
class ExecutionResult(Enum):
@@ -427,15 +431,7 @@ def _is_intermediate_output(dynprompt, node_id):
return getattr(class_def, 'HAS_INTERMEDIATE_OUTPUT', False)
def _send_cached_ui(server, node_id, display_node_id, cached, prompt_id, ui_outputs):
if cached.ui is not None:
ui_outputs[node_id] = cached.ui
if server.client_id is None:
return
cached_ui = cached.ui or {}
server.send_sync("executed", { "node": node_id, "display_node": display_node_id, "output": cached_ui.get("output", None), "prompt_id": prompt_id }, server.client_id)
async def execute(server, dynprompt, caches, current_item, extra_data, executed, prompt_id, execution_list, pending_subgraph_results, pending_async_nodes, ui_outputs):
async def execute(server: "ExecutionServer", dynprompt, caches, current_item, extra_data, executed, prompt_id, execution_list, pending_subgraph_results, pending_async_nodes, ui_outputs, asset_manager: AssetManager):
unique_id = current_item
real_node_id = dynprompt.get_real_node_id(unique_id)
display_node_id = dynprompt.get_display_node_id(unique_id)
@@ -445,7 +441,7 @@ async def execute(server, dynprompt, caches, current_item, extra_data, executed,
class_def = nodes.NODE_CLASS_MAPPINGS[class_type]
cached = await caches.outputs.get(unique_id)
if cached is not None:
_send_cached_ui(server, unique_id, display_node_id, cached, prompt_id, ui_outputs)
emit_cached_output(server, unique_id, display_node_id, cached, prompt_id, ui_outputs, asset_manager)
get_progress_state().finish_progress(unique_id)
execution_list.cache_update(unique_id, cached)
return (ExecutionResult.SUCCESS, None, None)
@@ -560,22 +556,19 @@ async def execute(server, dynprompt, caches, current_item, extra_data, executed,
unblock()
asyncio.create_task(await_completion())
return (ExecutionResult.PENDING, None, None)
cache_ui_value = ui_outputs.get(unique_id)
if len(output_ui) > 0:
# Enrich at output-processing time (not in the send path) so assets
# are registered even when no client is connected, and the asset id
# flows into ui_outputs and the cache alongside the raw entries.
output_ui = enrich_output_with_assets(output_ui)
ui_outputs[unique_id] = {
"meta": {
"node_id": unique_id,
"display_node": display_node_id,
"parent_node": parent_node_id,
"real_node_id": real_node_id,
},
"output": output_ui
meta = {
"node_id": unique_id,
"display_node": display_node_id,
"parent_node": parent_node_id,
"real_node_id": real_node_id,
}
enriched_output_ui = register_executed_outputs(output_ui, prompt_id, asset_manager)
ui_outputs[unique_id] = {"meta": meta, "output": enriched_output_ui}
cache_ui_value = {"meta": meta, "output": output_ui}
if server.client_id is not None:
server.send_sync("executed", { "node": unique_id, "display_node": display_node_id, "output": output_ui, "prompt_id": prompt_id }, server.client_id)
server.send_sync("executed", { "node": unique_id, "display_node": display_node_id, "output": enriched_output_ui, "prompt_id": prompt_id }, server.client_id)
if has_subgraph:
cached_outputs = []
new_node_ids = []
@@ -612,7 +605,7 @@ async def execute(server, dynprompt, caches, current_item, extra_data, executed,
pending_subgraph_results[unique_id] = cached_outputs
return (ExecutionResult.PENDING, None, None)
cache_entry = CacheEntry(ui=ui_outputs.get(unique_id), outputs=output_data)
cache_entry = CacheEntry(ui=cache_ui_value, outputs=output_data)
execution_list.cache_update(unique_id, cache_entry)
await caches.outputs.set(unique_id, cache_entry)
@@ -662,10 +655,11 @@ async def execute(server, dynprompt, caches, current_item, extra_data, executed,
return (ExecutionResult.SUCCESS, None, None)
class PromptExecutor:
def __init__(self, server, cache_type=False, cache_args=None):
def __init__(self, server: "ExecutionServer", cache_type=False, cache_args=None, asset_manager: AssetManager | None = None):
self.cache_args = cache_args
self.cache_type = cache_type
self.server = server
self.asset_manager = asset_manager if asset_manager is not None else default_asset_manager()
self.prompt_model_tracker = comfy.model_patcher.PromptModelTracker()
self.reset()
@@ -786,7 +780,7 @@ class PromptExecutor:
break
assert node_id is not None, "Node ID should not be None at this point"
result, error, ex = await execute(self.server, dynamic_prompt, self.caches, node_id, extra_data, executed, prompt_id, execution_list, pending_subgraph_results, pending_async_nodes, ui_node_outputs)
result, error, ex = await execute(self.server, dynamic_prompt, self.caches, node_id, extra_data, executed, prompt_id, execution_list, pending_subgraph_results, pending_async_nodes, ui_node_outputs, self.asset_manager)
self.success = result != ExecutionResult.FAILURE
if result == ExecutionResult.FAILURE:
self.handle_execution_error(prompt_id, dynamic_prompt.original_prompt, current_outputs, executed, error, ex)
@@ -820,7 +814,7 @@ class PromptExecutor:
cached = await self.caches.outputs.get(node_id)
if cached is not None:
display_node_id = dynamic_prompt.get_display_node_id(node_id)
_send_cached_ui(self.server, node_id, display_node_id, cached, prompt_id, ui_node_outputs)
emit_cached_output(self.server, node_id, display_node_id, cached, prompt_id, ui_node_outputs, self.asset_manager)
self.add_message("execution_success", { "prompt_id": prompt_id }, broadcast=False)
ui_outputs = {}
@@ -1249,7 +1243,7 @@ async def validate_prompt(prompt_id, prompt, partial_execution_list: Union[list[
MAXIMUM_HISTORY_SIZE = 10000
class PromptQueue:
def __init__(self, server):
def __init__(self, server: "ExecutionServer"):
self.server = server
self.mutex = threading.RLock()
self.not_empty = threading.Condition(self.mutex)

94
main.py
View File

@@ -22,8 +22,9 @@ console_log_level = get_console_log_level(args.verbose)
file_log_outputs = get_file_log_outputs(args.verbose)
setup_logger(log_level=console_log_level, file_outputs=file_log_outputs, use_stdout=args.log_stdout)
from app.assets.seeder import asset_seeder
from app.assets.services import register_output_files
from app.database.db import dependencies_available, init_db
from app.assets.lifecycle import cleanup_temp_filesystem
from app.assets.manager import AssetManager, default_asset_manager
import itertools
import utils.extra_config
from utils.mime_types import init_mime_types
@@ -34,7 +35,6 @@ import sys
from comfy_execution.progress import get_progress_state
from comfy_execution.utils import get_executing_context
from comfy_api import feature_flags
from app.database.db import init_db, dependencies_available
if __name__ == "__main__":
#NOTE: These do not do anything on core ComfyUI, they are for custom nodes.
@@ -316,39 +316,7 @@ def cuda_malloc_warning():
logging.warning("\nWARNING: this card most likely does not support cuda-malloc, if you get \"CUDA error\" please run ComfyUI with: --disable-cuda-malloc\n")
def _collect_output_absolute_paths(history_result: dict) -> list[str]:
"""Extract absolute file paths for output items from a history result."""
paths: list[str] = []
seen: set[str] = set()
for node_output in history_result.get("outputs", {}).values():
for items in node_output.values():
if not isinstance(items, list):
continue
for item in items:
if not isinstance(item, dict):
continue
item_type = item.get("type")
if item_type not in ("output", "temp"):
continue
base_dir = folder_paths.get_directory_by_type(item_type)
if base_dir is None:
continue
base_dir = os.path.abspath(base_dir)
filename = item.get("filename")
if not filename:
continue
abs_path = os.path.abspath(
os.path.join(base_dir, item.get("subfolder", ""), filename)
)
if not abs_path.startswith(base_dir + os.sep) and abs_path != base_dir:
continue
if abs_path not in seen:
seen.add(abs_path)
paths.append(abs_path)
return paths
def prompt_worker(q, server_instance):
def prompt_worker(q, server_instance, asset_manager):
current_time: float = 0.0
cache_ram = 0
cache_ram_inactive = 0
@@ -368,7 +336,7 @@ def prompt_worker(q, server_instance):
elif args.cache_none:
cache_type = execution.CacheType.NONE
e = execution.PromptExecutor(server_instance, cache_type=cache_type, cache_args={ "lru" : args.cache_lru, "ram" : cache_ram, "ram_inactive" : cache_ram_inactive } )
e = execution.PromptExecutor(server_instance, cache_type=cache_type, cache_args={ "lru" : args.cache_lru, "ram" : cache_ram, "ram_inactive" : cache_ram_inactive }, asset_manager=asset_manager )
last_gc_collect = 0
need_gc = False
gc_collect_interval = 10.0
@@ -390,7 +358,7 @@ def prompt_worker(q, server_instance):
for k in sensitive:
extra_data[k] = sensitive[k]
asset_seeder.pause()
asset_manager.pause_background_scan()
e.execute(item[2], prompt_id, extra_data, item[4])
need_gc = True
@@ -415,10 +383,6 @@ def prompt_worker(q, server_instance):
else:
logging.info("Prompt executed in {:.2f} seconds".format(execution_time), extra={'color': 'green'})
if not asset_seeder.is_disabled():
paths = _collect_output_absolute_paths(e.history_result)
register_output_files(paths, job_id=prompt_id)
flags = q.get_flags()
free_memory = flags.get("free_memory", False)
@@ -441,9 +405,8 @@ def prompt_worker(q, server_instance):
need_gc = False
hook_breaker_ac10a0.restore_functions()
if not asset_seeder.is_disabled():
asset_seeder.enqueue_enrich(roots=("output",), compute_hashes=args.enable_asset_hashing)
asset_seeder.resume()
asset_manager.queue_output_scan()
asset_manager.resume_background_scan()
async def run(server_instance, address='', port=8188, verbose=True, call_on_start=None):
@@ -486,19 +449,13 @@ def hijack_progress(server_instance):
comfy.utils.set_progress_bar_global_hook(hook)
def cleanup_temp():
temp_dir = folder_paths.get_temp_directory()
if os.path.exists(temp_dir):
shutil.rmtree(temp_dir, ignore_errors=True)
def setup_database(asset_manager):
if not dependencies_available():
return
def setup_database():
try:
if dependencies_available():
init_db()
if args.enable_assets:
if asset_seeder.start(roots=("models", "input", "output"), prune_first=True, compute_hashes=args.enable_asset_hashing):
logging.info("Background asset scan initiated for models, input, output")
init_db()
asset_manager.startup()
except Exception as e:
if "database is locked" in str(e):
logging.error(
@@ -507,6 +464,15 @@ def setup_database():
" --database-url sqlite:///path/to/another.db"
)
sys.exit(1)
if "Could not acquire lock on database" in str(e):
logging.error(
"Database is locked. Another ComfyUI process is already using this database.\n"
"To resolve this, specify a separate database file for this instance:\n"
" --database-url sqlite:///path/to/another.db"
)
if args.enable_assets:
sys.exit(1)
return
if args.enable_assets:
logging.error(
f"Failed to initialize database: {e}\n"
@@ -529,12 +495,15 @@ def start_comfyui(asyncio_loop=None):
temp_dir = os.path.join(os.path.abspath(args.temp_directory), "temp")
logging.info(f"Setting temp directory to: {temp_dir}")
folder_paths.set_temp_directory(temp_dir)
cleanup_temp()
asset_manager: AssetManager = default_asset_manager()
if not asset_manager.enabled:
cleanup_temp_filesystem()
if not asyncio_loop:
asyncio_loop = asyncio.new_event_loop()
asyncio.set_event_loop(asyncio_loop)
prompt_server = server.PromptServer(asyncio_loop)
prompt_server = server.PromptServer(asyncio_loop, asset_manager)
if args.enable_manager and not args.disable_manager_ui:
comfyui_manager.start()
@@ -552,12 +521,12 @@ def start_comfyui(asyncio_loop=None):
hook_breaker_ac10a0.restore_functions()
cuda_malloc_warning()
setup_database()
setup_database(asset_manager)
prompt_server.add_routes()
hijack_progress(prompt_server)
threading.Thread(target=prompt_worker, daemon=True, args=(prompt_server.prompt_queue, prompt_server,)).start()
threading.Thread(target=prompt_worker, daemon=True, args=(prompt_server.prompt_queue, prompt_server, asset_manager)).start()
if args.quick_test_for_ci:
exit(0)
@@ -605,7 +574,7 @@ if __name__ == "__main__":
"dynamic vram enabled and using native ComfyUI model formats instead. "
"ComfyUI native formats like fp8, int8 and w4a8 will be faster even if they are larger than your memory."
)
event_loop, _, start_all_func = start_comfyui()
event_loop, prompt_server, start_all_func = start_comfyui()
try:
x = start_all_func()
app.logger.print_startup_warnings()
@@ -613,5 +582,4 @@ if __name__ == "__main__":
except KeyboardInterrupt:
logging.info("\nStopped server")
finally:
asset_seeder.shutdown()
cleanup_temp()
prompt_server.asset_manager.shutdown()

View File

@@ -11,10 +11,6 @@ components:
description: Display name of the asset. Mirrors name for backwards compatibility.
nullable: true
type: string
file_path:
description: Relative path in global-namespace-root form (e.g. "models/checkpoints/flux.safetensors")
nullable: true
type: string
hash:
description: Blake3 hash of the asset content.
pattern: ^blake3:[a-f0-9]{64}$
@@ -152,10 +148,6 @@ components:
description: Display name of the asset. Mirrors name for backwards compatibility.
nullable: true
type: string
file_path:
description: Relative path in global-namespace-root form (e.g. "models/checkpoints/flux.safetensors")
nullable: true
type: string
hash:
description: Blake3 hash of the asset content.
pattern: ^blake3:[a-f0-9]{64}$
@@ -1714,11 +1706,6 @@ paths:
name: name_contains
schema:
type: string
- description: JSON object for filtering by metadata fields
in: query
name: metadata_filter
schema:
type: string
- description: Maximum number of assets to return (1-500)
in: query
name: limit
@@ -2552,11 +2539,6 @@ paths:
name: name_contains
schema:
type: string
- description: JSON object for filtering by metadata fields
in: query
name: metadata_filter
schema:
type: string
- description: Maximum number of tags to return (1-1000, default 100)
in: query
name: limit

View File

@@ -44,10 +44,6 @@ import node_helpers
from comfyui_version import __version__
from app.frontend_management import FrontendManager, parse_version
from comfy_api.internal import _ComfyNodeInternal
from app.assets.seeder import asset_seeder
from app.assets.api.routes import register_assets_routes
from app.assets.services.ingest import register_file_in_place
from app.assets.services.path_utils import get_known_subfolder_tags
from app.assets.services.asset_management import resolve_hash_to_path
from app.user_manager import UserManager
@@ -55,10 +51,13 @@ from app.model_manager import ModelFileManager
from app.custom_node_manager import CustomNodeManager
from app.subgraph_manager import SubgraphManager
from app.node_replace_manager import NodeReplaceManager
from typing import Optional, Union
from typing import TYPE_CHECKING, Optional, Union
from api_server.routes.internal.internal_routes import InternalRoutes
from protocol import BinaryEventTypes
if TYPE_CHECKING:
from app.assets.manager import AssetManager
# Import cache control middleware
from middleware.cache_middleware import cache_control
@@ -213,9 +212,10 @@ def create_block_external_middleware():
class PromptServer():
def __init__(self, loop):
def __init__(self, loop, asset_manager: "AssetManager"):
PromptServer.instance = self
self.asset_manager = asset_manager
self.user_manager = UserManager()
self.model_file_manager = ModelFileManager()
self.custom_node_manager = CustomNodeManager()
@@ -254,11 +254,8 @@ class PromptServer():
else args.front_end_root
)
logging.info(f"[Prompt Server] web root: {self.web_root}")
if args.enable_assets:
register_assets_routes(self.app, self.user_manager)
else:
register_assets_routes(self.app)
asset_seeder.disable()
self.asset_manager.register_routes(self.app, self.user_manager)
self.asset_manager.set_event_sink(self.send_sync)
routes = web.RouteTableDef()
self.routes = routes
self.last_node_id = None
@@ -440,22 +437,22 @@ class PromptServer():
resp = {"name" : filename, "subfolder": subfolder, "type": image_upload_type}
if args.enable_assets:
try:
tag = image_upload_type if image_upload_type in ("input", "output") else "input"
tags = [tag]
tags.extend(get_known_subfolder_tags(subfolder))
result = register_file_in_place(abs_path=filepath, name=filename, tags=tags)
resp["asset"] = {
"id": result.ref.id,
"name": result.ref.name,
"asset_hash": result.asset.hash,
"size": result.asset.size_bytes,
"mime_type": result.asset.mime_type,
"tags": result.tags,
}
except Exception:
logging.warning("Failed to register uploaded image as asset", exc_info=True)
view = self.asset_manager.register_upload(
abs_path=filepath,
name=filename,
upload_type=image_upload_type,
subfolder=subfolder,
content_written=not image_is_duplicate,
)
if view is not None:
resp["asset"] = {
"id": view.asset.id,
"name": view.asset.name,
"asset_hash": view.asset_hash,
"size": view.size,
"mime_type": view.mime_type,
"tags": view.tags,
}
return web.json_response(resp)
else:
@@ -523,8 +520,11 @@ class PromptServer():
# node preview, it constructs /view?filename=<asset_hash>, so this
# endpoint must resolve blake3 hashes to their on-disk file paths.
if filename.startswith("blake3:"):
owner_id = self.user_manager.get_request_user_id(request)
result = resolve_hash_to_path(filename, owner_id=owner_id)
# Side-effect call: get_request_user_id raises KeyError for an unknown or
# system user in multi-user mode, which is what gates hash resolution.
# The returned id is deliberately unused (resolution is not owner-scoped).
self.user_manager.get_request_user_id(request)
result = resolve_hash_to_path(filename)
if result is None:
return web.Response(status=404)
file, filename, resolved_content_type = result.abs_path, result.download_name, result.content_type
@@ -799,7 +799,7 @@ class PromptServer():
@routes.get("/object_info")
async def get_object_info(request):
asset_seeder.start(roots=("models", "input", "output"))
self.asset_manager.ensure_scan_started()
with folder_paths.cache_helper:
out = {}
for x in nodes.NODE_CLASS_MAPPINGS:

View File

@@ -0,0 +1,155 @@
import logging
import os
import sqlite3
import pytest
import torch
from alembic import command
from alembic.config import Config
from filelock import FileLock, Timeout
from app.database import db as db_module
from comfy.cli_args import args as cli_args
if not torch.cuda.is_available():
cli_args.cpu = True
import main # noqa: E402
_PRE_HEAD = "0006_add_loader_path"
def _make_config(db_path: str) -> Config:
root = os.path.join(os.path.dirname(__file__), "../..")
cfg = Config(os.path.abspath(os.path.join(root, "alembic.ini")))
cfg.set_main_option("script_location", os.path.abspath(os.path.join(root, "alembic_db")))
cfg.set_main_option("sqlalchemy.url", f"sqlite:///{db_path}")
return cfg
def _current_revision(db_path: str) -> str:
with sqlite3.connect(db_path) as conn:
rows = conn.execute("SELECT version_num FROM alembic_version").fetchall()
assert len(rows) == 1
return rows[0][0]
@pytest.fixture
def stale_db(tmp_path, monkeypatch):
db_path = str(tmp_path / "comfyui.db")
command.upgrade(_make_config(db_path), _PRE_HEAD)
monkeypatch.setattr(db_module.args, "database_url", f"sqlite:///{db_path}")
monkeypatch.setattr(db_module, "Session", None)
monkeypatch.setattr(db_module, "_db_lock", None)
yield db_path
if db_module._db_lock is not None:
db_module._db_lock.release(force=True)
def test_init_file_db_migrates_when_lock_is_free(stale_db):
db_module._init_file_db(db_module.args.database_url)
assert _current_revision(stale_db) != _PRE_HEAD
assert os.path.exists(stale_db + ".bkp")
def test_successful_init_keeps_holding_the_lock(stale_db):
db_module._init_file_db(db_module.args.database_url)
contender = FileLock(stale_db + ".lock")
with pytest.raises(Timeout):
contender.acquire(timeout=0)
def test_failed_init_releases_the_lock(stale_db, monkeypatch):
def _explode():
raise RuntimeError("alembic config exploded")
monkeypatch.setattr(db_module, "get_alembic_config", _explode)
with pytest.raises(RuntimeError, match="alembic config exploded"):
db_module._init_file_db(db_module.args.database_url)
contender = FileLock(stale_db + ".lock")
try:
contender.acquire(timeout=0)
except Timeout:
pytest.fail(
"a failed init stranded the lock; setup_database logs and CONTINUES when assets are "
"disabled, so this process would block every other instance for its whole lifetime "
"over a database it never opened"
)
contender.release()
def test_held_lock_blocks_before_any_migration_work(stale_db):
holder = FileLock(stale_db + ".lock")
holder.acquire(timeout=0)
try:
with pytest.raises(RuntimeError, match="Another ComfyUI process may already be using it"):
db_module._init_file_db(db_module.args.database_url)
assert not os.path.exists(stale_db + ".bkp")
assert _current_revision(stale_db) == _PRE_HEAD
assert db_module.Session is None
finally:
holder.release()
def test_setup_database_routes_file_lock_to_lock_guidance(monkeypatch, caplog):
monkeypatch.setattr(main, "dependencies_available", lambda: True)
def _raise_file_lock():
raise RuntimeError(
"Could not acquire lock on database 'x.db'. "
"Another ComfyUI process may already be using it. "
"Use --database-url to specify a separate database file."
)
monkeypatch.setattr(main, "init_db", _raise_file_lock)
monkeypatch.setattr(main.args, "enable_assets", False)
with caplog.at_level(logging.ERROR):
result = main.setup_database(None)
assert result is None
assert "Database is locked. Another ComfyUI process is already using this database." in caplog.text
assert "Failed to initialize database." not in caplog.text
def test_setup_database_exits_for_file_lock_when_assets_are_enabled(monkeypatch, caplog):
monkeypatch.setattr(main, "dependencies_available", lambda: True)
def _raise_file_lock():
raise RuntimeError(
"Could not acquire lock on database 'x.db'. "
"Another ComfyUI process may already be using it. "
"Use --database-url to specify a separate database file."
)
monkeypatch.setattr(main, "init_db", _raise_file_lock)
monkeypatch.setattr(main.args, "enable_assets", True)
with caplog.at_level(logging.ERROR), pytest.raises(SystemExit) as error:
main.setup_database(None)
assert error.value.code == 1
assert "Database is locked. Another ComfyUI process is already using this database." in caplog.text
assert "The --enable-assets flag requires a working database connection." not in caplog.text
def test_setup_database_exits_for_driver_lock_when_assets_are_disabled(monkeypatch, caplog):
monkeypatch.setattr(main, "dependencies_available", lambda: True)
def _raise_driver_lock():
raise Exception("sqlite3.OperationalError: database is locked")
monkeypatch.setattr(main, "init_db", _raise_driver_lock)
monkeypatch.setattr(main.args, "enable_assets", False)
with caplog.at_level(logging.ERROR), pytest.raises(SystemExit) as error:
main.setup_database(None)
assert error.value.code == 1
assert "Database is locked. Another ComfyUI process is already using this database." in caplog.text

View File

@@ -0,0 +1,215 @@
import os
import sqlite3
import pytest
from alembic import command
from alembic.config import Config
_BASELINE_0006 = "0006_add_loader_path"
def _make_config(db_path: str) -> Config:
root = os.path.join(os.path.dirname(__file__), "../..")
cfg = Config(os.path.abspath(os.path.join(root, "alembic.ini")))
cfg.set_main_option("script_location", os.path.abspath(os.path.join(root, "alembic_db")))
cfg.set_main_option("sqlalchemy.url", f"sqlite:///{db_path}")
return cfg
@pytest.fixture
def db_at_0006(tmp_path):
db_path = str(tmp_path / "test.db")
cfg = _make_config(db_path)
command.upgrade(cfg, _BASELINE_0006)
yield cfg, db_path
def test_0007_upgrade_from_0006(db_at_0006):
cfg, db_path = db_at_0006
command.upgrade(cfg, "head")
with sqlite3.connect(db_path) as conn:
tables = {r[0] for r in conn.execute("SELECT name FROM sqlite_master WHERE type='table'")}
assert "assets" in tables
assert "asset_contents" in tables
assert "asset_system_state" in tables
assert "asset_references" not in tables
def test_0007_schema_has_expected_columns(db_at_0006):
cfg, db_path = db_at_0006
command.upgrade(cfg, "head")
with sqlite3.connect(db_path) as conn:
asset_cols = {r[1] for r in conn.execute("PRAGMA table_info(assets)")}
content_cols = {r[1] for r in conn.execute("PRAGMA table_info(asset_contents)")}
assert {"id", "content_id", "name", "loader_path", "updated_at", "last_access_time"} <= asset_cols
assert {"id", "hash", "path", "is_missing", "mtime_ns"} <= content_cols
def test_0007_downgrade_restores_0006_schema(db_at_0006):
cfg, db_path = db_at_0006
command.upgrade(cfg, "head")
command.downgrade(cfg, _BASELINE_0006)
with sqlite3.connect(db_path) as conn:
tables = {r[0] for r in conn.execute("SELECT name FROM sqlite_master WHERE type='table'")}
assert "asset_references" in tables
assert "asset_contents" not in tables
def test_0007_invariants_on_migrated_db(db_at_0006):
cfg, db_path = db_at_0006
command.upgrade(cfg, "head")
with sqlite3.connect(db_path) as conn:
conn.execute("PRAGMA foreign_keys = ON")
conn.execute(
"INSERT INTO asset_contents(id, path, is_missing, size_bytes, created_at) "
"VALUES ('c1', '/tmp/f1', 0, 0, '2024-01-01')"
)
conn.execute(
"INSERT INTO assets(id, content_id, name, created_at, updated_at) "
"VALUES ('a1', 'c1', 'test', '2024-01-01', '2024-01-01')"
)
conn.commit()
with pytest.raises(sqlite3.IntegrityError):
conn.execute(
"INSERT INTO asset_contents(id, path, is_missing, size_bytes, created_at) "
"VALUES ('c2', '/tmp/f1', 0, 0, '2024-01-01')"
)
conn.commit()
conn.rollback()
conn.execute(
"INSERT INTO asset_contents(id, path, is_missing, size_bytes, created_at) "
"VALUES ('c3', '/tmp/f1', 1, 0, '2024-01-01')"
)
conn.commit()
conn.execute("UPDATE asset_contents SET hash='abc123' WHERE id='c1'")
conn.execute(
"INSERT INTO asset_contents(id, path, hash, is_missing, size_bytes, created_at) "
"VALUES ('c4', '/tmp/f2', 'abc123', 0, 0, '2024-01-01')"
)
conn.commit()
def test_0007_orm_parity(db_at_0006, tmp_path):
from sqlalchemy import create_engine, inspect
import app.assets.database.models as asset_models
from app.database.models import Base
cfg, db_path = db_at_0006
command.upgrade(cfg, "head")
with sqlite3.connect(db_path) as conn:
alembic_tables = {
r[0]
for r in conn.execute(
"SELECT name FROM sqlite_master WHERE type='table' "
"AND name NOT LIKE 'alembic%' AND name NOT LIKE 'sqlite%'"
)
}
alembic_engine = create_engine(f"sqlite:///{db_path}")
orm_db = str(tmp_path / "orm.db")
engine = create_engine(f"sqlite:///{orm_db}")
alembic_inspector = inspect(alembic_engine)
Base.metadata.create_all(engine)
orm_inspector = inspect(engine)
orm_tables = set(orm_inspector.get_table_names())
assert asset_models.AssetMeta.__tablename__ == "asset_meta"
alembic_indexes = {
(index["name"], tuple(index["column_names"]))
for index in alembic_inspector.get_indexes("asset_meta")
}
orm_indexes = {
(index.name, tuple(index.columns.keys()))
for index in Base.metadata.tables["asset_meta"].indexes
}
assert alembic_tables == orm_tables, f"Mismatch: alembic={alembic_tables}, orm={orm_tables}"
assert alembic_indexes == orm_indexes, f"Mismatch: alembic={alembic_indexes}, orm={orm_indexes}"
def test_0007_downgrade_chain_past_0003_succeeds(db_at_0006):
cfg, db_path = db_at_0006
command.upgrade(cfg, "head")
command.downgrade(cfg, "0001_assets")
with sqlite3.connect(db_path) as conn:
tables = {r[0] for r in conn.execute("SELECT name FROM sqlite_master WHERE type='table'")}
assert "assets_info" in tables
assert "asset_references" not in tables
assert "asset_contents" not in tables
def test_0007_downgrade_restores_legacy_assets_definition(db_at_0006, tmp_path):
from sqlalchemy import create_engine, inspect
cfg, db_path = db_at_0006
reference_db = str(tmp_path / "reference_0006.db")
reference_cfg = _make_config(reference_db)
command.upgrade(reference_cfg, _BASELINE_0006)
command.upgrade(cfg, "head")
command.downgrade(cfg, _BASELINE_0006)
def _assets_shape(path: str):
engine = create_engine(f"sqlite:///{path}")
try:
inspector = inspect(engine)
indexes = {
(index["name"], tuple(index["column_names"]), bool(index["unique"]))
for index in inspector.get_indexes("assets")
}
checks = {
(check["name"], check["sqltext"])
for check in inspector.get_check_constraints("assets")
}
defaults = {
column["name"]: column["default"]
for column in inspector.get_columns("assets")
}
return indexes, checks, defaults
finally:
engine.dispose()
restored_indexes, restored_checks, restored_defaults = _assets_shape(db_path)
expected_indexes, expected_checks, expected_defaults = _assets_shape(reference_db)
assert restored_indexes == expected_indexes, (
f"index drift: restored={restored_indexes}, expected={expected_indexes}"
)
assert restored_checks == expected_checks, (
f"check drift: restored={restored_checks}, expected={expected_checks}"
)
assert restored_defaults == expected_defaults, (
f"default drift: restored={restored_defaults}, expected={expected_defaults}"
)
def test_0007_downgrade_restores_0006_asset_references_indexes(db_at_0006, tmp_path):
from sqlalchemy import create_engine, inspect
cfg, db_path = db_at_0006
reference_db = str(tmp_path / "reference_0006.db")
reference_cfg = _make_config(reference_db)
command.upgrade(reference_cfg, _BASELINE_0006)
command.upgrade(cfg, "head")
command.downgrade(cfg, _BASELINE_0006)
def _asset_reference_indexes(path: str) -> set[tuple[str, tuple[str, ...]]]:
engine = create_engine(f"sqlite:///{path}")
try:
return {
(index["name"], tuple(index["column_names"]))
for index in inspect(engine).get_indexes("asset_references")
}
finally:
engine.dispose()
restored = _asset_reference_indexes(db_path)
expected = _asset_reference_indexes(reference_db)
assert restored == expected, f"index drift: restored={restored}, expected={expected}"

View File

@@ -0,0 +1,56 @@
import os
import sqlite3
import pytest
from alembic import command
from alembic.config import Config
def _make_config(db_path: str) -> Config:
root = os.path.join(os.path.dirname(__file__), "../..")
cfg = Config(os.path.abspath(os.path.join(root, "alembic.ini")))
cfg.set_main_option("script_location", os.path.abspath(os.path.join(root, "alembic_db")))
cfg.set_main_option("sqlalchemy.url", f"sqlite:///{db_path}")
return cfg
def _user_tables(db_path: str) -> set[str]:
conn = sqlite3.connect(db_path)
try:
return {
row[0]
for row in conn.execute(
"SELECT name FROM sqlite_master WHERE type='table' "
"AND name NOT LIKE 'alembic%' AND name NOT LIKE 'sqlite%'"
)
}
finally:
conn.close()
@pytest.fixture
def db_at_head(tmp_path):
db_path = str(tmp_path / "roundtrip.db")
cfg = _make_config(db_path)
command.upgrade(cfg, "head")
yield cfg, db_path
def test_downgrade_to_base_empties_the_schema(db_at_head):
cfg, db_path = db_at_head
assert _user_tables(db_path), "fixture did not build a schema"
command.downgrade(cfg, "base")
assert _user_tables(db_path) == set()
def test_upgrade_to_head_again_after_downgrade_to_base(db_at_head):
cfg, db_path = db_at_head
expected = _user_tables(db_path)
command.downgrade(cfg, "base")
command.upgrade(cfg, "head")
assert _user_tables(db_path) == expected

View File

@@ -0,0 +1,110 @@
import os
import pytest
from alembic import command
from alembic.config import Config
from alembic.script import ScriptDirectory
import app.database.db as db_module
_BASELINE_0006 = "0006_add_loader_path"
_BASELINE_0002 = "0002_merge_to_asset_references"
_DESTRUCTIVE = "0007_record_content_split"
_DISCARDED_CLASSES = (
"manual tags",
"user metadata",
"previews",
"renames",
"API-created records",
"job_id",
)
def _make_config(db_path: str) -> Config:
root = os.path.join(os.path.dirname(__file__), "../..")
cfg = Config(os.path.abspath(os.path.join(root, "alembic.ini")))
cfg.set_main_option("script_location", os.path.abspath(os.path.join(root, "alembic_db")))
cfg.set_main_option("sqlalchemy.url", f"sqlite:///{db_path}")
return cfg
@pytest.fixture
def captured_warnings(monkeypatch):
warnings: list[str] = []
monkeypatch.setattr(db_module, "log_startup_warning", warnings.append)
return warnings
@pytest.fixture
def db_url_for(monkeypatch):
def _set(db_path: str) -> str:
url = f"sqlite:///{db_path}"
monkeypatch.setattr(db_module.args, "database_url", url)
return url
return _set
def _migrate(db_path: str, db_url: str, db_exists: bool) -> None:
db_module._migrate_and_bind(db_url, db_path, db_exists)
def test_upgrade_across_the_destructive_revision_warns_and_names_the_backup(
tmp_path, captured_warnings, db_url_for
):
db_path = str(tmp_path / "existing.db")
command.upgrade(_make_config(db_path), _BASELINE_0006)
db_url = db_url_for(db_path)
assert os.path.exists(db_path)
_migrate(db_path, db_url, db_exists=True)
assert len(captured_warnings) == 1, (
"the destructive rebuild is one startup event, not one warning per traversed revision"
)
message = captured_warnings[0]
assert db_path + ".bkp" in message, "the warning must name the backup it left behind"
for discarded in _DISCARDED_CLASSES:
assert discarded in message, (
f"the warning must inventory what the rebuild discarded; {discarded!r} is missing"
)
def test_fresh_install_warns_nothing(tmp_path, captured_warnings, db_url_for):
db_path = str(tmp_path / "fresh.db")
db_url = db_url_for(db_path)
assert not os.path.exists(db_path)
_migrate(db_path, db_url, db_exists=False)
assert captured_warnings == [], (
"a fresh install discarded nothing and took no backup, so it has nothing to warn about"
)
def test_upgrade_that_does_not_cross_the_destructive_revision_warns_nothing(
tmp_path, captured_warnings, db_url_for, monkeypatch
):
db_path = str(tmp_path / "partial.db")
command.upgrade(_make_config(db_path), _BASELINE_0002)
db_url = db_url_for(db_path)
monkeypatch.setattr(ScriptDirectory, "get_current_head", lambda _self: _BASELINE_0006)
_migrate(db_path, db_url, db_exists=True)
assert captured_warnings == [], (
"a multi-revision upgrade that never traverses the rebuild keeps every row it had"
)
def test_database_already_at_head_warns_nothing(tmp_path, captured_warnings, db_url_for):
db_path = str(tmp_path / "current.db")
command.upgrade(_make_config(db_path), "head")
db_url = db_url_for(db_path)
_migrate(db_path, db_url, db_exists=True)
assert captured_warnings == [], (
"an up-to-date database migrates nothing, so it must stay silent on every boot"
)

View File

View File

@@ -13,6 +13,8 @@ from typing import Callable, Iterator, Optional
import pytest
import requests
from .helpers import assert_hash_fields_consistent
def pytest_addoption(parser: pytest.Parser) -> None:
"""
@@ -28,6 +30,18 @@ def pytest_addoption(parser: pytest.Parser) -> None:
default=os.environ.get("ASSETS_TEST_DB_URL"),
help="SQLAlchemy DB URL (e.g. sqlite:///path/to/db.sqlite3)",
)
parser.addoption(
"--enable-asset-hashing",
action="store_true",
help="Start the assets subprocess with hash-mode behavior enabled.",
)
def pytest_configure(config: pytest.Config) -> None:
config.addinivalue_line(
"markers",
"hashing_on: exercises the subprocess harness with --enable-asset-hashing",
)
def _free_port() -> int:
@@ -103,8 +117,7 @@ def comfy_url_and_proc(comfy_tmp_base_dir: Path, request: pytest.FixtureRequest)
if not (comfy_root / "main.py").is_file():
raise FileNotFoundError(f"main.py not found under {comfy_root}")
proc = subprocess.Popen(
args=[
command = [
sys.executable,
"main.py",
f"--base-directory={str(comfy_tmp_base_dir)}",
@@ -115,7 +128,19 @@ def comfy_url_and_proc(comfy_tmp_base_dir: Path, request: pytest.FixtureRequest)
"--port",
str(port),
"--cpu",
],
]
if (
request.config.getoption("--enable-asset-hashing")
or "hashing_on" in request.config.getoption("markexpr")
or any(
item.get_closest_marker("hashing_on")
for item in request.session.items
)
):
command.append("--enable-asset-hashing")
proc = subprocess.Popen(
args=command,
stdout=out_log,
stderr=err_log,
cwd=str(comfy_root),
@@ -190,8 +215,9 @@ def _post_multipart_asset(
@pytest.fixture
def make_asset_bytes() -> Callable[[str, int], bytes]:
# Salt content per test so it never collides with assets left over from
# earlier tests. Delete is now always a soft delete (content is preserved),
# so the suite can no longer rely on hard-deleting content for isolation.
# earlier tests. Delete hard-deletes the record but preserves content
# (content rows and files are untouched), so the suite cannot rely on delete
# removing content for isolation.
# Deterministic within a test: the same (name, size) yields the same bytes.
salt = uuid.uuid4().bytes
@@ -236,9 +262,6 @@ def seeded_asset(request: pytest.FixtureRequest, http: requests.Session, api_bas
if tags is None:
tags = ["models", "model_type:checkpoints", "unit-tests", "alpha"]
meta = {"purpose": "test", "epoch": 1, "flags": ["x", "y"], "nullable": None}
# Unique content per test so the seed always creates a fresh asset (201).
# Delete is now always a soft delete, so content from a prior test survives
# and would otherwise dedup this upload into an existing asset (200).
content = uuid.uuid4().bytes + b"A" * (4096 - 16)
files = {"file": (name, content, "application/octet-stream")}
form_data = {
@@ -249,7 +272,6 @@ def seeded_asset(request: pytest.FixtureRequest, http: requests.Session, api_bas
r = http.post(api_base + "/api/assets", files=files, data=form_data, timeout=120)
body = r.json()
assert r.status_code == 201, body
from helpers import assert_hash_fields_consistent
assert_hash_fields_consistent(body)
return body

View File

@@ -1,7 +1,124 @@
"""Helper functions for assets integration tests."""
import time
from __future__ import annotations
import json
import time
import uuid
from collections.abc import Iterator, Mapping
from dataclasses import dataclass
from datetime import datetime, timedelta
from typing import NotRequired, TypeAlias, TypedDict
import pytest
import requests
from aiohttp import web
from aiohttp.test_utils import make_mocked_request
from sqlalchemy import Engine, create_engine
from sqlalchemy.orm import Session
from app.assets.api import routes
from app.assets.database.models import Asset
from app.assets.database.queries.records import create_content, create_record
from app.database.models import Base
class AssetItem(TypedDict):
id: str
name: str
preview_id: NotRequired[str]
class AssetListBody(TypedDict):
assets: list[AssetItem]
total: int
has_more: bool
next_cursor: NotRequired[str]
class ErrorItem(TypedDict):
code: str
class ErrorBody(TypedDict):
error: ErrorItem
@dataclass(frozen=True, slots=True)
class RecordSeed:
name: str
tags: tuple[str, ...] = ()
size_bytes: int = 0
RouteDatabase: TypeAlias = tuple[Engine, Session]
@pytest.fixture(autouse=True)
def autoclean_unit_test_assets() -> Iterator[None]:
yield
@pytest.fixture
def route_database(monkeypatch: pytest.MonkeyPatch) -> Iterator[RouteDatabase]:
engine = create_engine("sqlite:///:memory:")
Base.metadata.create_all(engine)
monkeypatch.setattr(routes, "create_session", lambda: Session(engine))
monkeypatch.setattr(routes, "_ASSETS_ENABLED", True)
with Session(engine) as session:
yield engine, session
engine.dispose()
def seed_record(session: Session, seed: RecordSeed) -> Asset:
content = create_content(
session,
path=f"/output/{uuid.uuid4()}-{seed.name}",
size_bytes=seed.size_bytes,
)
return create_record(
session,
content_id=content.id,
name=seed.name,
mime_type="image/png",
tags=seed.tags,
)
@pytest.fixture
def sortable_record_ids(route_database: RouteDatabase) -> tuple[str, str]:
_, session = route_database
older = seed_record(session, RecordSeed("z.png", ("sort-case",), 100))
newer = seed_record(session, RecordSeed("a.png", ("sort-case",), 200))
base_time = datetime(2026, 1, 1)
older.created_at = base_time
newer.created_at = base_time + timedelta(days=1)
newer.updated_at = base_time
older.updated_at = base_time + timedelta(days=1)
older.last_access_time = base_time
newer.last_access_time = base_time + timedelta(days=1)
session.commit()
return newer.id, older.id
async def request_assets(query: str = "") -> web.StreamResponse:
suffix = f"?{query}" if query else ""
return await routes.list_assets_route(
make_mocked_request("GET", f"/api/assets{suffix}")
)
def asset_list_body(response: web.StreamResponse) -> AssetListBody:
assert isinstance(response, web.Response)
body = response.body
assert isinstance(body, bytes | bytearray)
return json.loads(body)
def error_body(response: web.StreamResponse) -> ErrorBody:
assert isinstance(response, web.Response)
body = response.body
assert isinstance(body, bytes | bytearray)
return json.loads(body)
def trigger_sync_seed_assets(session: requests.Session, base_url: str) -> None:
@@ -28,7 +145,10 @@ def get_asset_filename(asset_hash: str, extension: str) -> str:
return asset_hash.removeprefix("blake3:") + extension
def assert_hash_fields_consistent(body: dict, expected_hash: str | None = None) -> None:
def assert_hash_fields_consistent(
body: Mapping[str, str | None],
expected_hash: str | None = None,
) -> None:
"""Assert hash and asset_hash invariants on an Asset response.
Both must be present or both absent (so a regression that drops only one

View File

@@ -5,6 +5,18 @@ from sqlalchemy.orm import Session
from app.assets.database.models import Base
@pytest.fixture(scope="session", autouse=True)
def assert_asset_metadata_tables():
assert set(Base.metadata.tables) == {
"assets",
"asset_contents",
"asset_meta",
"asset_tags",
"tags",
"asset_system_state",
}
@pytest.fixture
def session():
"""In-memory SQLite session for fast unit tests."""

View File

@@ -1,187 +1,20 @@
import uuid
import pytest
import sqlalchemy as sa
from sqlalchemy.orm import Session
from app.assets.helpers import get_utc_now
from app.assets.database.models import Asset
from app.assets.database.queries import (
asset_exists_by_hash,
get_asset_by_hash,
upsert_asset,
bulk_insert_assets,
update_asset_hash_and_mime,
)
from app.assets.database.models import AssetContent
from app.assets.database.queries import create_content, create_record
class TestAssetExistsByHash:
@pytest.mark.parametrize(
"setup_hash,query_hash,expected",
[
(None, "nonexistent", False), # No asset exists
("blake3:abc123", "blake3:abc123", True), # Asset exists with matching hash
(None, "", False), # Null hash in DB doesn't match empty string
],
ids=["nonexistent", "existing", "null_hash_no_match"],
)
def test_exists_by_hash(self, session: Session, setup_hash, query_hash, expected):
if setup_hash is not None or query_hash == "":
asset = Asset(hash=setup_hash, size_bytes=100)
session.add(asset)
session.commit()
def test_hash_is_a_content_fact(session: Session) -> None:
content = create_content(session, "/models/a.safetensors", hash="same", size_bytes=12)
record = create_record(session, content.id, "a.safetensors")
assert asset_exists_by_hash(session, asset_hash=query_hash) is expected
assert session.scalar(sa.select(AssetContent).where(AssetContent.hash == "same")) == content
assert record.content_id == content.id
class TestGetAssetByHash:
@pytest.mark.parametrize(
"setup_hash,query_hash,should_find",
[
(None, "nonexistent", False),
("blake3:def456", "blake3:def456", True),
],
ids=["nonexistent", "existing"],
)
def test_get_by_hash(self, session: Session, setup_hash, query_hash, should_find):
if setup_hash is not None:
asset = Asset(hash=setup_hash, size_bytes=200, mime_type="image/png")
session.add(asset)
session.commit()
def test_equal_hashes_do_not_merge_content_rows(session: Session) -> None:
first = create_content(session, "/models/a.safetensors", hash="same")
second = create_content(session, "/models/b.safetensors", hash="same")
result = get_asset_by_hash(session, asset_hash=query_hash)
if should_find:
assert result is not None
assert result.size_bytes == 200
assert result.mime_type == "image/png"
else:
assert result is None
class TestUpsertAsset:
@pytest.mark.parametrize(
"first_size,first_mime,second_size,second_mime,expect_created,expect_updated,final_size,final_mime",
[
# New asset creation
(None, None, 1024, "application/octet-stream", True, False, 1024, "application/octet-stream"),
# Existing asset, same values - no update
(500, "text/plain", 500, "text/plain", False, False, 500, "text/plain"),
# Existing asset with size 0, update with new values
(0, None, 2048, "image/png", False, True, 2048, "image/png"),
# Existing asset, second call with size 0 - no update
(1000, None, 0, None, False, False, 1000, None),
],
ids=["new_asset", "existing_no_change", "update_from_zero", "zero_size_no_update"],
)
def test_upsert_scenarios(
self,
session: Session,
first_size,
first_mime,
second_size,
second_mime,
expect_created,
expect_updated,
final_size,
final_mime,
):
asset_hash = f"blake3:test_{first_size}_{second_size}"
# First upsert (if first_size is not None, we're testing the second call)
if first_size is not None:
upsert_asset(
session,
asset_hash=asset_hash,
size_bytes=first_size,
mime_type=first_mime,
)
session.commit()
# The upsert call we're testing
asset, created, updated = upsert_asset(
session,
asset_hash=asset_hash,
size_bytes=second_size,
mime_type=second_mime,
)
session.commit()
assert created is expect_created
assert updated is expect_updated
assert asset.size_bytes == final_size
assert asset.mime_type == final_mime
class TestBulkInsertAssets:
def test_inserts_multiple_assets(self, session: Session):
now = get_utc_now()
rows = [
{"id": str(uuid.uuid4()), "hash": "blake3:bulk1", "size_bytes": 100, "mime_type": "text/plain", "created_at": now},
{"id": str(uuid.uuid4()), "hash": "blake3:bulk2", "size_bytes": 200, "mime_type": "image/png", "created_at": now},
{"id": str(uuid.uuid4()), "hash": "blake3:bulk3", "size_bytes": 300, "mime_type": None, "created_at": now},
]
bulk_insert_assets(session, rows)
session.commit()
assets = session.query(Asset).all()
assert len(assets) == 3
hashes = {a.hash for a in assets}
assert hashes == {"blake3:bulk1", "blake3:bulk2", "blake3:bulk3"}
def test_empty_list_is_noop(self, session: Session):
bulk_insert_assets(session, [])
session.commit()
assert session.query(Asset).count() == 0
def test_handles_large_batch(self, session: Session):
"""Test chunking logic with more rows than MAX_BIND_PARAMS allows."""
now = get_utc_now()
rows = [
{"id": str(uuid.uuid4()), "hash": f"blake3:large{i}", "size_bytes": i, "mime_type": None, "created_at": now}
for i in range(200)
]
bulk_insert_assets(session, rows)
session.commit()
assert session.query(Asset).count() == 200
class TestMimeTypeImmutability:
"""mime_type on Asset is write-once: set on first ingest, never overwritten."""
@pytest.mark.parametrize(
"initial_mime,second_mime,expected_mime",
[
("image/png", "image/jpeg", "image/png"),
(None, "image/png", "image/png"),
],
ids=["preserves_existing", "fills_null"],
)
def test_upsert_mime_immutability(self, session: Session, initial_mime, second_mime, expected_mime):
h = f"blake3:upsert_{initial_mime}_{second_mime}"
upsert_asset(session, asset_hash=h, size_bytes=100, mime_type=initial_mime)
session.commit()
asset, created, _ = upsert_asset(session, asset_hash=h, size_bytes=100, mime_type=second_mime)
assert created is False
assert asset.mime_type == expected_mime
@pytest.mark.parametrize(
"initial_mime,update_mime,update_hash,expected_mime,expected_hash",
[
(None, "image/png", None, "image/png", "blake3:upd0"),
("image/png", "image/jpeg", None, "image/png", "blake3:upd1"),
("image/png", "image/jpeg", "blake3:upd2_new", "image/png", "blake3:upd2_new"),
],
ids=["fills_null", "preserves_existing", "hash_updates_mime_locked"],
)
def test_update_asset_hash_and_mime_immutability(
self, session: Session, initial_mime, update_mime, update_hash, expected_mime, expected_hash,
):
h = expected_hash.removesuffix("_new")
asset = Asset(hash=h, size_bytes=100, mime_type=initial_mime)
session.add(asset)
session.flush()
update_asset_hash_and_mime(session, asset_id=asset.id, mime_type=update_mime, asset_hash=update_hash)
assert asset.mime_type == expected_mime
assert asset.hash == expected_hash
assert first.id != second.id

View File

@@ -1,580 +1,29 @@
import time
import uuid
import pytest
from sqlalchemy.orm import Session
from app.assets.database.models import Asset, AssetReference, AssetReferenceMeta
from app.assets.database.queries import (
reference_exists_for_asset_id,
get_reference_by_id,
insert_reference,
get_or_create_reference,
update_reference_timestamps,
list_references_page,
fetch_reference_asset_and_tags,
fetch_reference_and_asset,
update_reference_access_time,
set_reference_metadata,
delete_reference_by_id,
set_reference_preview,
bulk_insert_references_ignore_conflicts,
get_reference_ids_by_ids,
ensure_tags_exist,
add_tags_to_reference,
create_content,
create_record,
delete_record,
get_record_by_id,
rename_record,
)
from app.assets.helpers import get_utc_now
def _make_asset(session: Session, hash_val: str | None = None, size: int = 1024) -> Asset:
asset = Asset(hash=hash_val, size_bytes=size, mime_type="application/octet-stream")
session.add(asset)
session.flush()
return asset
def test_record_operations_do_not_mutate_content(session: Session) -> None:
content = create_content(session, "/output/result.png", hash="digest", size_bytes=3)
record = create_record(session, content.id, "before.png", mime_type="image/png")
renamed = rename_record(session, record.id, "after.png")
def _make_reference(
session: Session,
asset: Asset,
name: str = "test",
owner_id: str = "",
) -> AssetReference:
now = get_utc_now()
ref = AssetReference(
owner_id=owner_id,
name=name,
asset_id=asset.id,
created_at=now,
updated_at=now,
last_access_time=now,
)
session.add(ref)
session.flush()
return ref
assert renamed.content_id == content.id
assert renamed.name == "after.png"
assert content.hash == "digest"
class TestReferenceExistsForAssetId:
def test_returns_false_when_no_reference(self, session: Session):
asset = _make_asset(session, "hash1")
assert reference_exists_for_asset_id(session, asset_id=asset.id) is False
def test_delete_record_is_hard_delete(session: Session) -> None:
content = create_content(session, "/output/result.png")
record = create_record(session, content.id, "result.png")
def test_returns_true_when_reference_exists(self, session: Session):
asset = _make_asset(session, "hash1")
_make_reference(session, asset)
assert reference_exists_for_asset_id(session, asset_id=asset.id) is True
delete_record(session, record.id)
class TestGetReferenceById:
def test_returns_none_for_nonexistent(self, session: Session):
assert get_reference_by_id(session, reference_id="nonexistent") is None
def test_returns_reference(self, session: Session):
asset = _make_asset(session, "hash1")
ref = _make_reference(session, asset, name="myfile.txt")
result = get_reference_by_id(session, reference_id=ref.id)
assert result is not None
assert result.name == "myfile.txt"
class TestListReferencesPage:
def test_empty_db(self, session: Session):
refs, tag_map, total = list_references_page(session)
assert refs == []
assert tag_map == {}
assert total == 0
def test_returns_references_with_tags(self, session: Session):
asset = _make_asset(session, "hash1")
ref = _make_reference(session, asset, name="test.bin")
ensure_tags_exist(session, ["alpha", "beta"])
add_tags_to_reference(session, reference_id=ref.id, tags=["alpha", "beta"])
session.commit()
refs, tag_map, total = list_references_page(session)
assert len(refs) == 1
assert refs[0].id == ref.id
assert set(tag_map[ref.id]) == {"alpha", "beta"}
assert total == 1
def test_name_contains_filter(self, session: Session):
asset = _make_asset(session, "hash1")
_make_reference(session, asset, name="model_v1.safetensors")
_make_reference(session, asset, name="config.json")
session.commit()
refs, _, total = list_references_page(session, name_contains="model")
assert total == 1
assert refs[0].name == "model_v1.safetensors"
def test_owner_visibility(self, session: Session):
asset = _make_asset(session, "hash1")
_make_reference(session, asset, name="public", owner_id="")
_make_reference(session, asset, name="private", owner_id="user1")
session.commit()
# Empty owner sees only public
refs, _, total = list_references_page(session, owner_id="")
assert total == 1
assert refs[0].name == "public"
# Owner sees both
refs, _, total = list_references_page(session, owner_id="user1")
assert total == 2
def test_include_tags_filter(self, session: Session):
asset = _make_asset(session, "hash1")
ref1 = _make_reference(session, asset, name="tagged")
_make_reference(session, asset, name="untagged")
ensure_tags_exist(session, ["wanted"])
add_tags_to_reference(session, reference_id=ref1.id, tags=["wanted"])
session.commit()
refs, _, total = list_references_page(session, include_tags=["wanted"])
assert total == 1
assert refs[0].name == "tagged"
def test_include_tags_filter_ands_persisted_model_tags(self, session: Session):
asset = _make_asset(session, "hash-model-tags")
checkpoint = _make_reference(session, asset, name="checkpoint")
lora = _make_reference(session, asset, name="lora")
input_ref = _make_reference(session, asset, name="input")
ensure_tags_exist(
session,
["models", "model_type:checkpoints", "model_type:loras", "unit-tests"],
)
add_tags_to_reference(
session,
reference_id=checkpoint.id,
tags=["models", "model_type:checkpoints", "unit-tests"],
origin="automatic",
)
add_tags_to_reference(
session,
reference_id=lora.id,
tags=["models", "model_type:loras", "unit-tests"],
origin="automatic",
)
add_tags_to_reference(
session,
reference_id=input_ref.id,
tags=["unit-tests"],
)
session.commit()
refs, _, total = list_references_page(
session,
include_tags=["models", "model_type:checkpoints", "unit-tests"],
)
assert total == 1
assert refs[0].id == checkpoint.id
def test_include_tags_filter_preserves_model_type_case(self, session: Session):
asset = _make_asset(session, "hash-model-case")
ref = _make_reference(session, asset, name="llm")
ensure_tags_exist(session, ["models", "model_type:LLM"])
add_tags_to_reference(
session,
reference_id=ref.id,
tags=["models", "model_type:LLM"],
origin="automatic",
)
session.commit()
refs, _, total = list_references_page(
session, include_tags=["models", "model_type:LLM"]
)
refs_lower, _, total_lower = list_references_page(
session, include_tags=["models", "model_type:llm"]
)
assert total == 1
assert refs[0].id == ref.id
assert total_lower == 0
assert refs_lower == []
def test_exclude_tags_filter(self, session: Session):
asset = _make_asset(session, "hash1")
_make_reference(session, asset, name="keep")
ref_exclude = _make_reference(session, asset, name="exclude")
ensure_tags_exist(session, ["bad"])
add_tags_to_reference(session, reference_id=ref_exclude.id, tags=["bad"])
session.commit()
refs, _, total = list_references_page(session, exclude_tags=["bad"])
assert total == 1
assert refs[0].name == "keep"
def test_sorting(self, session: Session):
asset = _make_asset(session, "hash1", size=100)
asset2 = _make_asset(session, "hash2", size=500)
_make_reference(session, asset, name="small")
_make_reference(session, asset2, name="large")
session.commit()
refs, _, _ = list_references_page(session, sort="size", order="desc")
assert refs[0].name == "large"
refs, _, _ = list_references_page(session, sort="name", order="asc")
assert refs[0].name == "large"
class TestFetchReferenceAssetAndTags:
def test_returns_none_for_nonexistent(self, session: Session):
result = fetch_reference_asset_and_tags(session, "nonexistent")
assert result is None
def test_returns_tuple(self, session: Session):
asset = _make_asset(session, "hash1")
ref = _make_reference(session, asset, name="test.bin")
ensure_tags_exist(session, ["tag1"])
add_tags_to_reference(session, reference_id=ref.id, tags=["tag1"])
session.commit()
result = fetch_reference_asset_and_tags(session, ref.id)
assert result is not None
ret_ref, ret_asset, ret_tags = result
assert ret_ref.id == ref.id
assert ret_asset.id == asset.id
assert ret_tags == ["tag1"]
class TestFetchReferenceAndAsset:
def test_returns_none_for_nonexistent(self, session: Session):
result = fetch_reference_and_asset(session, reference_id="nonexistent")
assert result is None
def test_returns_tuple(self, session: Session):
asset = _make_asset(session, "hash1")
ref = _make_reference(session, asset)
session.commit()
result = fetch_reference_and_asset(session, reference_id=ref.id)
assert result is not None
ret_ref, ret_asset = result
assert ret_ref.id == ref.id
assert ret_asset.id == asset.id
class TestUpdateReferenceAccessTime:
def test_updates_last_access_time(self, session: Session):
asset = _make_asset(session, "hash1")
ref = _make_reference(session, asset)
original_time = ref.last_access_time
session.commit()
import time
time.sleep(0.01)
update_reference_access_time(session, reference_id=ref.id)
session.commit()
session.refresh(ref)
assert ref.last_access_time > original_time
class TestDeleteReferenceById:
def test_deletes_existing(self, session: Session):
asset = _make_asset(session, "hash1")
ref = _make_reference(session, asset)
session.commit()
result = delete_reference_by_id(session, reference_id=ref.id, owner_id="")
assert result is True
assert get_reference_by_id(session, reference_id=ref.id) is None
def test_returns_false_for_nonexistent(self, session: Session):
result = delete_reference_by_id(session, reference_id="nonexistent", owner_id="")
assert result is False
def test_respects_owner_visibility(self, session: Session):
asset = _make_asset(session, "hash1")
ref = _make_reference(session, asset, owner_id="user1")
session.commit()
result = delete_reference_by_id(session, reference_id=ref.id, owner_id="user2")
assert result is False
assert get_reference_by_id(session, reference_id=ref.id) is not None
class TestSetReferencePreview:
def test_sets_preview(self, session: Session):
asset = _make_asset(session, "hash1")
preview_asset = _make_asset(session, "preview_hash")
ref = _make_reference(session, asset)
preview_ref = _make_reference(session, preview_asset, name="preview.png")
session.commit()
set_reference_preview(session, reference_id=ref.id, preview_reference_id=preview_ref.id)
session.commit()
session.refresh(ref)
assert ref.preview_id == preview_ref.id
def test_clears_preview(self, session: Session):
asset = _make_asset(session, "hash1")
preview_asset = _make_asset(session, "preview_hash")
ref = _make_reference(session, asset)
preview_ref = _make_reference(session, preview_asset, name="preview.png")
ref.preview_id = preview_ref.id
session.commit()
set_reference_preview(session, reference_id=ref.id, preview_reference_id=None)
session.commit()
session.refresh(ref)
assert ref.preview_id is None
def test_raises_for_nonexistent_reference(self, session: Session):
with pytest.raises(ValueError, match="not found"):
set_reference_preview(session, reference_id="nonexistent", preview_reference_id=None)
def test_raises_for_nonexistent_preview(self, session: Session):
asset = _make_asset(session, "hash1")
ref = _make_reference(session, asset)
session.commit()
with pytest.raises(ValueError, match="Preview AssetReference"):
set_reference_preview(session, reference_id=ref.id, preview_reference_id="nonexistent")
class TestInsertReference:
def test_creates_new_reference(self, session: Session):
asset = _make_asset(session, "hash1")
ref = insert_reference(
session, asset_id=asset.id, owner_id="user1", name="test.bin"
)
session.commit()
assert ref is not None
assert ref.name == "test.bin"
assert ref.owner_id == "user1"
def test_allows_duplicate_names(self, session: Session):
asset = _make_asset(session, "hash1")
ref1 = insert_reference(session, asset_id=asset.id, owner_id="user1", name="dup.bin")
session.commit()
# Duplicate names are now allowed
ref2 = insert_reference(
session, asset_id=asset.id, owner_id="user1", name="dup.bin"
)
session.commit()
assert ref1 is not None
assert ref2 is not None
assert ref1.id != ref2.id
class TestGetOrCreateReference:
def test_creates_new_reference(self, session: Session):
asset = _make_asset(session, "hash1")
ref, created = get_or_create_reference(
session, asset_id=asset.id, owner_id="user1", name="new.bin"
)
session.commit()
assert created is True
assert ref.name == "new.bin"
def test_always_creates_new_reference(self, session: Session):
asset = _make_asset(session, "hash1")
ref1, created1 = get_or_create_reference(
session, asset_id=asset.id, owner_id="user1", name="existing.bin"
)
session.commit()
# Duplicate names are allowed, so always creates new
ref2, created2 = get_or_create_reference(
session, asset_id=asset.id, owner_id="user1", name="existing.bin"
)
session.commit()
assert created1 is True
assert created2 is True
assert ref1.id != ref2.id
class TestUpdateReferenceTimestamps:
def test_updates_timestamps(self, session: Session):
asset = _make_asset(session, "hash1")
ref = _make_reference(session, asset)
original_updated_at = ref.updated_at
session.commit()
time.sleep(0.01)
update_reference_timestamps(session, ref)
session.commit()
session.refresh(ref)
assert ref.updated_at > original_updated_at
def test_updates_preview_id(self, session: Session):
asset = _make_asset(session, "hash1")
preview_asset = _make_asset(session, "preview_hash")
ref = _make_reference(session, asset)
preview_ref = _make_reference(session, preview_asset, name="preview.png")
session.commit()
update_reference_timestamps(session, ref, preview_id=preview_ref.id)
session.commit()
session.refresh(ref)
assert ref.preview_id == preview_ref.id
class TestSetReferenceMetadata:
def test_sets_metadata(self, session: Session):
asset = _make_asset(session, "hash1")
ref = _make_reference(session, asset)
session.commit()
set_reference_metadata(
session, reference_id=ref.id, user_metadata={"key": "value"}
)
session.commit()
session.refresh(ref)
assert ref.user_metadata == {"key": "value"}
# Check metadata table
meta = session.query(AssetReferenceMeta).filter_by(asset_reference_id=ref.id).all()
assert len(meta) == 1
assert meta[0].key == "key"
assert meta[0].val_str == "value"
def test_replaces_existing_metadata(self, session: Session):
asset = _make_asset(session, "hash1")
ref = _make_reference(session, asset)
session.commit()
set_reference_metadata(
session, reference_id=ref.id, user_metadata={"old": "data"}
)
session.commit()
set_reference_metadata(
session, reference_id=ref.id, user_metadata={"new": "data"}
)
session.commit()
meta = session.query(AssetReferenceMeta).filter_by(asset_reference_id=ref.id).all()
assert len(meta) == 1
assert meta[0].key == "new"
def test_clears_metadata_with_empty_dict(self, session: Session):
asset = _make_asset(session, "hash1")
ref = _make_reference(session, asset)
session.commit()
set_reference_metadata(
session, reference_id=ref.id, user_metadata={"key": "value"}
)
session.commit()
set_reference_metadata(
session, reference_id=ref.id, user_metadata={}
)
session.commit()
session.refresh(ref)
assert ref.user_metadata == {}
meta = session.query(AssetReferenceMeta).filter_by(asset_reference_id=ref.id).all()
assert len(meta) == 0
def test_raises_for_nonexistent(self, session: Session):
with pytest.raises(ValueError, match="not found"):
set_reference_metadata(
session, reference_id="nonexistent", user_metadata={"key": "value"}
)
class TestBulkInsertReferencesIgnoreConflicts:
def test_inserts_multiple_references(self, session: Session):
asset = _make_asset(session, "hash1")
now = get_utc_now()
rows = [
{
"id": str(uuid.uuid4()),
"owner_id": "",
"name": "bulk1.bin",
"asset_id": asset.id,
"preview_id": None,
"user_metadata": {},
"created_at": now,
"updated_at": now,
"last_access_time": now,
},
{
"id": str(uuid.uuid4()),
"owner_id": "",
"name": "bulk2.bin",
"asset_id": asset.id,
"preview_id": None,
"user_metadata": {},
"created_at": now,
"updated_at": now,
"last_access_time": now,
},
]
bulk_insert_references_ignore_conflicts(session, rows)
session.commit()
refs = session.query(AssetReference).all()
assert len(refs) == 2
def test_allows_duplicate_names(self, session: Session):
asset = _make_asset(session, "hash1")
_make_reference(session, asset, name="existing.bin", owner_id="")
session.commit()
now = get_utc_now()
rows = [
{
"id": str(uuid.uuid4()),
"owner_id": "",
"name": "existing.bin",
"asset_id": asset.id,
"preview_id": None,
"user_metadata": {},
"created_at": now,
"updated_at": now,
"last_access_time": now,
},
{
"id": str(uuid.uuid4()),
"owner_id": "",
"name": "new.bin",
"asset_id": asset.id,
"preview_id": None,
"user_metadata": {},
"created_at": now,
"updated_at": now,
"last_access_time": now,
},
]
bulk_insert_references_ignore_conflicts(session, rows)
session.commit()
# Duplicate names allowed, so all 3 rows exist
refs = session.query(AssetReference).all()
assert len(refs) == 3
def test_empty_list_is_noop(self, session: Session):
bulk_insert_references_ignore_conflicts(session, [])
assert session.query(AssetReference).count() == 0
class TestGetReferenceIdsByIds:
def test_returns_existing_ids(self, session: Session):
asset = _make_asset(session, "hash1")
ref1 = _make_reference(session, asset, name="a.bin")
ref2 = _make_reference(session, asset, name="b.bin")
session.commit()
found = get_reference_ids_by_ids(session, [ref1.id, ref2.id, "nonexistent"])
assert found == {ref1.id, ref2.id}
def test_empty_list_returns_empty(self, session: Session):
found = get_reference_ids_by_ids(session, [])
assert found == set()
assert get_record_by_id(session, record.id) is None

View File

@@ -1,112 +1,33 @@
"""Keyset-pagination tiebreaker tests for list_references_page.
When multiple rows share the same primary sort value (e.g. four assets
created in the same microsecond), the secondary `ORDER BY id` is what keeps
keyset pagination from losing or repeating rows. This file exercises that
branch directly against an in-memory SQLite session — engineering identical
timestamps via HTTP is unreliable enough that we work at the query layer.
"""
import uuid
from datetime import datetime
import pytest
from sqlalchemy.orm import Session
from app.assets.database.models import Asset, AssetReference
from app.assets.database.queries.asset_reference import list_references_page
from app.assets.database.queries import create_content, create_record, list_records_page
from app.assets.database.queries.records import RecordCursorBoundary, RecordPageSpec
def _make_ref(session: Session, created_at: datetime, name: str, owner: str = "") -> AssetReference:
asset = Asset(hash=f"blake3:{uuid.uuid4().hex}", size_bytes=1024)
session.add(asset)
def test_record_keyset_cursor_pages_in_creation_order(session: Session) -> None:
records = [
create_record(session, create_content(session, f"/output/{name}").id, name)
for name in ("one.png", "two.png", "three.png")
]
for index, record in enumerate(records, start=1):
record.id = f"00000000-0000-0000-0000-{index:012d}"
session.flush()
ref = AssetReference(
id=str(uuid.uuid4()),
asset_id=asset.id,
owner_id=owner,
name=name,
file_path=f"/tmp/{name}",
created_at=created_at,
updated_at=created_at,
last_access_time=created_at,
is_missing=False,
first_page, _, _ = list_records_page(
session,
RecordPageSpec(limit=2, order="asc"),
)
session.add(ref)
return ref
@pytest.mark.parametrize("order", ["desc", "asc"])
def test_tiebreaker_walks_duplicate_sort_values(session: Session, order: str):
"""Four rows with the SAME created_at must paginate cleanly under cursor
mode — no row dropped, no row repeated, despite the primary sort column
being non-discriminating.
"""
shared_ts = datetime(2024, 5, 20, 12, 0, 0) # naive UTC, like the DB stores
refs = [_make_ref(session, shared_ts, f"tie_{i}.png") for i in range(4)]
session.commit()
expected_ids = sorted([r.id for r in refs], reverse=(order == "desc"))
# Walk the cursor by hand: page size 2, take 3 pages (2 + 2 + 0).
seen: list[str] = []
after_value = None
after_id = None
for _ in range(4): # generous loop bound; ought to be 2 iterations
page, _tag_map, _total = list_references_page(
session,
boundary_record = first_page[-1]
second_page, _, _ = list_records_page(
session,
RecordPageSpec(
limit=2,
sort="created_at",
order=order,
after_cursor_value=after_value,
after_cursor_id=after_id,
)
if not page:
break
seen.extend(p.id for p in page)
# Use the last row's (created_at, id) as the next cursor input.
last = page[-1]
after_value, after_id = last.created_at, last.id
if len(page) < 2:
break
assert seen == expected_ids, (
f"keyset tiebreaker failed for order={order}: expected {expected_ids}, got {seen}"
order="asc",
after=RecordCursorBoundary(
value=boundary_record.created_at,
id=boundary_record.id,
),
),
)
def test_tiebreaker_no_duplicates_under_mixed_collisions(session: Session):
"""Some rows share a timestamp, some don't. The cursor must still walk
every row exactly once regardless of where ties sit relative to a
page boundary."""
t1 = datetime(2024, 5, 20, 12, 0, 0)
t2 = datetime(2024, 5, 20, 12, 0, 1)
layout = [t1, t1, t1, t2, t2] # three rows at t1, two at t2
refs = [_make_ref(session, ts, f"mix_{i}.png") for i, ts in enumerate(layout)]
session.commit()
all_ids = {r.id for r in refs}
seen_set: set[str] = set()
seen_list: list[str] = []
after_value = None
after_id = None
for _ in range(6):
page, _, _ = list_references_page(
session,
limit=2,
sort="created_at",
order="desc",
after_cursor_value=after_value,
after_cursor_id=after_id,
)
if not page:
break
for p in page:
assert p.id not in seen_set, f"duplicate row {p.id} appeared in cursor walk"
seen_set.add(p.id)
seen_list.append(p.id)
last = page[-1]
after_value, after_id = last.created_at, last.id
if len(page) < 2:
break
assert seen_set == all_ids, f"missing rows: expected {all_ids}, got {seen_set}"
assert [record.id for record in first_page + second_page] == [record.id for record in records]

View File

@@ -1,532 +1,32 @@
"""Tests for cache_state (AssetReference file path) query functions."""
import pytest
from sqlalchemy.orm import Session
from app.assets.database.models import Asset, AssetReference
from app.assets.database.queries import (
list_references_by_asset_id,
upsert_reference,
get_unreferenced_unhashed_asset_ids,
delete_assets_by_ids,
get_references_for_prefixes,
bulk_update_needs_verify,
delete_references_by_ids,
delete_orphaned_seed_asset,
bulk_insert_references_ignore_conflicts,
get_references_by_paths_and_asset_ids,
mark_references_missing_outside_prefixes,
restore_references_by_paths,
create_content,
create_record,
fetch_record_tags,
mark_content_missing,
unset_content_missing,
)
from app.assets.helpers import select_best_live_path, get_utc_now
def _make_asset(session: Session, hash_val: str | None = None, size: int = 1024) -> Asset:
asset = Asset(hash=hash_val, size_bytes=size)
session.add(asset)
session.flush()
return asset
def test_missing_state_lives_on_content_and_projects_to_records(session: Session) -> None:
content = create_content(session, "/models/model.safetensors")
first = create_record(session, content.id, "model-a")
second = create_record(session, content.id, "model-b")
mark_content_missing(session, content.id)
def _make_reference(
session: Session,
asset: Asset,
file_path: str,
name: str = "test",
mtime_ns: int | None = None,
needs_verify: bool = False,
) -> AssetReference:
now = get_utc_now()
ref = AssetReference(
asset_id=asset.id,
file_path=file_path,
name=name,
mtime_ns=mtime_ns,
needs_verify=needs_verify,
created_at=now,
updated_at=now,
last_access_time=now,
)
session.add(ref)
session.flush()
return ref
assert content.is_missing is True
assert fetch_record_tags(session, first.id) == ["missing"]
assert fetch_record_tags(session, second.id) == ["missing"]
class TestListReferencesByAssetId:
def test_returns_empty_for_no_references(self, session: Session):
asset = _make_asset(session, "hash1")
refs = list_references_by_asset_id(session, asset_id=asset.id)
assert list(refs) == []
def test_recovery_clears_the_automatic_missing_marker(session: Session) -> None:
content = create_content(session, "/models/model.safetensors")
record = create_record(session, content.id, "model")
mark_content_missing(session, content.id)
def test_returns_references_for_asset(self, session: Session):
asset = _make_asset(session, "hash1")
_make_reference(session, asset, "/path/a.bin", name="a")
_make_reference(session, asset, "/path/b.bin", name="b")
session.commit()
unset_content_missing(session, content.id)
refs = list_references_by_asset_id(session, asset_id=asset.id)
paths = [r.file_path for r in refs]
assert set(paths) == {"/path/a.bin", "/path/b.bin"}
def test_does_not_return_other_assets_references(self, session: Session):
asset1 = _make_asset(session, "hash1")
asset2 = _make_asset(session, "hash2")
_make_reference(session, asset1, "/path/asset1.bin", name="a1")
_make_reference(session, asset2, "/path/asset2.bin", name="a2")
session.commit()
refs = list_references_by_asset_id(session, asset_id=asset1.id)
paths = [r.file_path for r in refs]
assert paths == ["/path/asset1.bin"]
class TestSelectBestLivePath:
def test_returns_empty_for_empty_list(self):
result = select_best_live_path([])
assert result == ""
def test_returns_empty_when_no_files_exist(self, session: Session):
asset = _make_asset(session, "hash1")
ref = _make_reference(session, asset, "/nonexistent/path.bin")
session.commit()
result = select_best_live_path([ref])
assert result == ""
def test_prefers_verified_path(self, session: Session, tmp_path):
"""needs_verify=False should be preferred."""
asset = _make_asset(session, "hash1")
verified_file = tmp_path / "verified.bin"
verified_file.write_bytes(b"data")
unverified_file = tmp_path / "unverified.bin"
unverified_file.write_bytes(b"data")
ref_verified = _make_reference(
session, asset, str(verified_file), name="verified", needs_verify=False
)
ref_unverified = _make_reference(
session, asset, str(unverified_file), name="unverified", needs_verify=True
)
session.commit()
refs = [ref_unverified, ref_verified]
result = select_best_live_path(refs)
assert result == str(verified_file)
def test_falls_back_to_existing_unverified(self, session: Session, tmp_path):
"""If all references need verification, return first existing path."""
asset = _make_asset(session, "hash1")
existing_file = tmp_path / "exists.bin"
existing_file.write_bytes(b"data")
ref = _make_reference(session, asset, str(existing_file), needs_verify=True)
session.commit()
result = select_best_live_path([ref])
assert result == str(existing_file)
class TestSelectBestLivePathWithMocking:
def test_handles_missing_file_path_attr(self):
"""Gracefully handle references with None file_path."""
class MockRef:
file_path = None
needs_verify = False
result = select_best_live_path([MockRef()])
assert result == ""
class TestUpsertReference:
@pytest.mark.parametrize(
"initial_mtime,second_mtime,expect_created,expect_updated,final_mtime",
[
# New reference creation
(None, 12345, True, False, 12345),
# Existing reference, same mtime - no update
(100, 100, False, False, 100),
# Existing reference, different mtime - update
(100, 200, False, True, 200),
],
ids=["new_reference", "existing_no_change", "existing_update_mtime"],
)
def test_upsert_scenarios(
self, session: Session, initial_mtime, second_mtime, expect_created, expect_updated, final_mtime
):
asset = _make_asset(session, "hash1")
file_path = f"/path_{initial_mtime}_{second_mtime}.bin"
name = f"file_{initial_mtime}_{second_mtime}"
# Create initial reference if needed
if initial_mtime is not None:
upsert_reference(session, asset_id=asset.id, file_path=file_path, name=name, mtime_ns=initial_mtime)
session.commit()
# The upsert call we're testing
created, updated = upsert_reference(
session, asset_id=asset.id, file_path=file_path, name=name, mtime_ns=second_mtime
)
session.commit()
assert created is expect_created
assert updated is expect_updated
ref = session.query(AssetReference).filter_by(file_path=file_path).one()
assert ref.mtime_ns == final_mtime
def test_upsert_refreshes_loader_path_on_existing_reference(self, session: Session):
"""Re-ingesting an existing reference writes the loader_path computed
by that ingest, healing NULL or stale values even when nothing else
about the row changed."""
asset = _make_asset(session, "hash1")
file_path = "/models/checkpoints/sub/model.safetensors"
upsert_reference(
session, asset_id=asset.id, file_path=file_path, name="model",
mtime_ns=100, loader_path=None,
)
session.commit()
created, updated = upsert_reference(
session, asset_id=asset.id, file_path=file_path, name="model",
mtime_ns=100, loader_path="sub/model.safetensors",
)
session.commit()
assert created is False
assert updated is True
ref = session.query(AssetReference).filter_by(file_path=file_path).one()
assert ref.loader_path == "sub/model.safetensors"
# Identical loader_path is a no-op, not a spurious update.
created, updated = upsert_reference(
session, asset_id=asset.id, file_path=file_path, name="model",
mtime_ns=100, loader_path="sub/model.safetensors",
)
session.commit()
assert created is False
assert updated is False
def test_upsert_restores_missing_reference(self, session: Session):
"""Upserting a reference that was marked missing should restore it."""
asset = _make_asset(session, "hash1")
file_path = "/restored/file.bin"
ref = _make_reference(session, asset, file_path, mtime_ns=100)
ref.is_missing = True
session.commit()
created, updated = upsert_reference(
session, asset_id=asset.id, file_path=file_path, name="restored", mtime_ns=100
)
session.commit()
assert created is False
assert updated is True
restored_ref = session.query(AssetReference).filter_by(file_path=file_path).one()
assert restored_ref.is_missing is False
class TestRestoreReferencesByPaths:
def test_restores_missing_references(self, session: Session):
asset = _make_asset(session, "hash1")
missing_path = "/missing/file.bin"
active_path = "/active/file.bin"
missing_ref = _make_reference(session, asset, missing_path, name="missing")
missing_ref.is_missing = True
_make_reference(session, asset, active_path, name="active")
session.commit()
restored = restore_references_by_paths(session, [missing_path])
session.commit()
assert restored == 1
ref = session.query(AssetReference).filter_by(file_path=missing_path).one()
assert ref.is_missing is False
def test_empty_list_restores_nothing(self, session: Session):
restored = restore_references_by_paths(session, [])
assert restored == 0
class TestMarkReferencesMissingOutsidePrefixes:
def test_marks_references_missing_outside_prefixes(self, session: Session, tmp_path):
asset = _make_asset(session, "hash1")
valid_dir = tmp_path / "valid"
valid_dir.mkdir()
invalid_dir = tmp_path / "invalid"
invalid_dir.mkdir()
valid_path = str(valid_dir / "file.bin")
invalid_path = str(invalid_dir / "file.bin")
_make_reference(session, asset, valid_path, name="valid")
_make_reference(session, asset, invalid_path, name="invalid")
session.commit()
marked = mark_references_missing_outside_prefixes(session, [str(valid_dir)])
session.commit()
assert marked == 1
all_refs = session.query(AssetReference).all()
assert len(all_refs) == 2
valid_ref = next(r for r in all_refs if r.file_path == valid_path)
invalid_ref = next(r for r in all_refs if r.file_path == invalid_path)
assert valid_ref.is_missing is False
assert invalid_ref.is_missing is True
def test_empty_prefixes_marks_nothing(self, session: Session):
asset = _make_asset(session, "hash1")
_make_reference(session, asset, "/some/path.bin")
session.commit()
marked = mark_references_missing_outside_prefixes(session, [])
assert marked == 0
class TestGetUnreferencedUnhashedAssetIds:
def test_returns_unreferenced_unhashed_assets(self, session: Session):
# Unhashed asset (hash=None) with no references (no file_path)
no_refs = _make_asset(session, hash_val=None)
# Unhashed asset with active reference (not unreferenced)
with_active_ref = _make_asset(session, hash_val=None)
_make_reference(session, with_active_ref, "/has/ref.bin", name="has_ref")
# Unhashed asset with only missing reference (should be unreferenced)
with_missing_ref = _make_asset(session, hash_val=None)
missing_ref = _make_reference(session, with_missing_ref, "/missing/ref.bin", name="missing_ref")
missing_ref.is_missing = True
# Regular asset (hash not None) - should not be returned
_make_asset(session, hash_val="blake3:regular")
session.commit()
unreferenced = get_unreferenced_unhashed_asset_ids(session)
assert no_refs.id in unreferenced
assert with_missing_ref.id in unreferenced
assert with_active_ref.id not in unreferenced
class TestDeleteAssetsByIds:
def test_deletes_assets_and_references(self, session: Session):
asset = _make_asset(session, "hash1")
_make_reference(session, asset, "/test/path.bin", name="test")
session.commit()
deleted = delete_assets_by_ids(session, [asset.id])
session.commit()
assert deleted == 1
assert session.query(Asset).count() == 0
assert session.query(AssetReference).count() == 0
def test_empty_list_deletes_nothing(self, session: Session):
_make_asset(session, "hash1")
session.commit()
deleted = delete_assets_by_ids(session, [])
assert deleted == 0
assert session.query(Asset).count() == 1
class TestGetReferencesForPrefixes:
def test_returns_references_matching_prefix(self, session: Session, tmp_path):
asset = _make_asset(session, "hash1")
dir1 = tmp_path / "dir1"
dir1.mkdir()
dir2 = tmp_path / "dir2"
dir2.mkdir()
path1 = str(dir1 / "file.bin")
path2 = str(dir2 / "file.bin")
_make_reference(session, asset, path1, name="file1", mtime_ns=100)
_make_reference(session, asset, path2, name="file2", mtime_ns=200)
session.commit()
rows = get_references_for_prefixes(session, [str(dir1)])
assert len(rows) == 1
assert rows[0].file_path == path1
def test_empty_prefixes_returns_empty(self, session: Session):
asset = _make_asset(session, "hash1")
_make_reference(session, asset, "/some/path.bin")
session.commit()
rows = get_references_for_prefixes(session, [])
assert rows == []
class TestBulkSetNeedsVerify:
def test_sets_needs_verify_flag(self, session: Session):
asset = _make_asset(session, "hash1")
ref1 = _make_reference(session, asset, "/path1.bin", needs_verify=False)
ref2 = _make_reference(session, asset, "/path2.bin", needs_verify=False)
session.commit()
updated = bulk_update_needs_verify(session, [ref1.id, ref2.id], True)
session.commit()
assert updated == 2
session.refresh(ref1)
session.refresh(ref2)
assert ref1.needs_verify is True
assert ref2.needs_verify is True
def test_empty_list_updates_nothing(self, session: Session):
updated = bulk_update_needs_verify(session, [], True)
assert updated == 0
class TestDeleteReferencesByIds:
def test_deletes_references_by_id(self, session: Session):
asset = _make_asset(session, "hash1")
ref1 = _make_reference(session, asset, "/path1.bin")
_make_reference(session, asset, "/path2.bin")
session.commit()
deleted = delete_references_by_ids(session, [ref1.id])
session.commit()
assert deleted == 1
assert session.query(AssetReference).count() == 1
def test_empty_list_deletes_nothing(self, session: Session):
deleted = delete_references_by_ids(session, [])
assert deleted == 0
class TestDeleteOrphanedSeedAsset:
@pytest.mark.parametrize(
"create_asset,expected_deleted,expected_count",
[
(True, True, 0), # Existing asset gets deleted
(False, False, 0), # Nonexistent returns False
],
ids=["deletes_existing", "nonexistent_returns_false"],
)
def test_delete_orphaned_seed_asset(
self, session: Session, create_asset, expected_deleted, expected_count
):
asset_id = "nonexistent-id"
if create_asset:
asset = _make_asset(session, hash_val=None)
asset_id = asset.id
_make_reference(session, asset, "/test/path.bin", name="test")
session.commit()
deleted = delete_orphaned_seed_asset(session, asset_id)
if create_asset:
session.commit()
assert deleted is expected_deleted
assert session.query(Asset).count() == expected_count
class TestBulkInsertReferencesIgnoreConflicts:
def test_inserts_multiple_references(self, session: Session):
asset = _make_asset(session, "hash1")
now = get_utc_now()
rows = [
{
"asset_id": asset.id,
"file_path": "/bulk1.bin",
"name": "bulk1",
"mtime_ns": 100,
"created_at": now,
"updated_at": now,
"last_access_time": now,
},
{
"asset_id": asset.id,
"file_path": "/bulk2.bin",
"name": "bulk2",
"mtime_ns": 200,
"created_at": now,
"updated_at": now,
"last_access_time": now,
},
]
bulk_insert_references_ignore_conflicts(session, rows)
session.commit()
assert session.query(AssetReference).count() == 2
def test_ignores_conflicts(self, session: Session):
asset = _make_asset(session, "hash1")
_make_reference(session, asset, "/existing.bin", mtime_ns=100)
session.commit()
now = get_utc_now()
rows = [
{
"asset_id": asset.id,
"file_path": "/existing.bin",
"name": "existing",
"mtime_ns": 999,
"created_at": now,
"updated_at": now,
"last_access_time": now,
},
{
"asset_id": asset.id,
"file_path": "/new.bin",
"name": "new",
"mtime_ns": 200,
"created_at": now,
"updated_at": now,
"last_access_time": now,
},
]
bulk_insert_references_ignore_conflicts(session, rows)
session.commit()
assert session.query(AssetReference).count() == 2
existing = session.query(AssetReference).filter_by(file_path="/existing.bin").one()
assert existing.mtime_ns == 100 # Original value preserved
def test_empty_list_is_noop(self, session: Session):
bulk_insert_references_ignore_conflicts(session, [])
assert session.query(AssetReference).count() == 0
class TestGetReferencesByPathsAndAssetIds:
def test_returns_matching_paths(self, session: Session):
asset1 = _make_asset(session, "hash1")
asset2 = _make_asset(session, "hash2")
_make_reference(session, asset1, "/path1.bin")
_make_reference(session, asset2, "/path2.bin")
session.commit()
path_to_asset = {
"/path1.bin": asset1.id,
"/path2.bin": asset2.id,
}
winners = get_references_by_paths_and_asset_ids(session, path_to_asset)
assert winners == {"/path1.bin", "/path2.bin"}
def test_excludes_non_matching_asset_ids(self, session: Session):
asset1 = _make_asset(session, "hash1")
asset2 = _make_asset(session, "hash2")
_make_reference(session, asset1, "/path1.bin")
session.commit()
# Path exists but with different asset_id
path_to_asset = {"/path1.bin": asset2.id}
winners = get_references_by_paths_and_asset_ids(session, path_to_asset)
assert winners == set()
def test_empty_dict_returns_empty(self, session: Session):
winners = get_references_by_paths_and_asset_ids(session, {})
assert winners == set()
assert content.is_missing is False
assert fetch_record_tags(session, record.id) == []

View File

@@ -0,0 +1,171 @@
import os
from datetime import datetime
from unittest.mock import patch
import pytest
from sqlalchemy import create_engine, update
from sqlalchemy.orm import Session
import app.assets.mode as mode_module
from app.assets.database.models import AssetContent
from app.assets.database.queries.records import create_content
from app.assets.services.lookup import (
claim_qualified_content,
is_temp_path as _is_temp_path,
lookup_for_from_hash,
lookup_for_view,
refresh_qualified_content,
)
from app.database.models import Base
@pytest.fixture
def session(tmp_path):
engine = create_engine(
f"sqlite:///{tmp_path}/test.db", connect_args={"check_same_thread": False}
)
Base.metadata.create_all(engine)
with Session(engine) as sess:
yield sess
@pytest.fixture(autouse=True)
def enable_hashing():
class FakeArgs:
enable_asset_hashing = True
mode_module.init(FakeArgs())
yield
mode_module.init(None)
def _make_file(tmp_path, name: str, content: bytes = b"bytes") -> str:
p = tmp_path / name
p.write_bytes(content)
return str(p)
def test_temp_only_match_from_hash_returns_none(session, tmp_path):
f = _make_file(tmp_path, "f.png")
create_content(session, path=f, hash="abc123")
session.commit()
with patch("app.assets.services.lookup.is_temp_path", return_value=True):
result = lookup_for_from_hash(session, "abc123")
assert result is None
def test_temp_only_match_view_returns_none(session, tmp_path):
f = _make_file(tmp_path, "f2.png")
create_content(session, path=f, hash="abc123")
session.commit()
with patch("app.assets.services.lookup.is_temp_path", return_value=True):
result = lookup_for_view(session, "abc123")
assert result is None
def test_sibling_prefix_not_temp(tmp_path):
with patch("folder_paths.get_temp_directory", return_value=str(tmp_path / "temp")):
assert not _is_temp_path(str(tmp_path / "temp-other" / "f.png"))
assert _is_temp_path(str(tmp_path / "temp" / "f.png"))
def test_off_mode_from_hash_returns_none(session, tmp_path):
class FakeArgs:
enable_asset_hashing = False
mode_module.init(FakeArgs())
f = _make_file(tmp_path, "f3.png")
create_content(session, path=f, hash="abc123")
session.commit()
result = lookup_for_from_hash(session, "abc123")
assert result is None
def test_upload_content_lookup_not_gated_on_hashing_flag(session, tmp_path):
class FakeArgs:
enable_asset_hashing = False
mode_module.init(FakeArgs())
f = _make_file(tmp_path, "f4.png")
content = create_content(session, path=f, hash="abc123", size_bytes=os.path.getsize(f))
session.commit()
result = lookup_for_view(session, "abc123")
assert result is not None
assert result.id == content.id
def test_stale_older_newer_live_returns_newer(session, tmp_path):
f_newer = _make_file(tmp_path, "newer.png")
old_time = datetime(2020, 1, 1)
new_time = datetime(2024, 1, 1)
c_old = create_content(session, path="/nonexistent/old.png", hash="xyz")
c_new = create_content(session, path=f_newer, hash="xyz", size_bytes=os.path.getsize(f_newer))
session.execute(update(AssetContent).where(AssetContent.id == c_old.id).values(created_at=old_time))
session.execute(update(AssetContent).where(AssetContent.id == c_new.id).values(created_at=new_time))
session.commit()
result = lookup_for_from_hash(session, "xyz")
assert result is not None
assert result.id == c_new.id
def test_claim_qualified_content_true_for_a_live_matching_row(session, tmp_path):
f = _make_file(tmp_path, "revalidate_ok.png")
content = create_content(session, path=f, hash="dup")
session.commit()
assert claim_qualified_content(session, content.id, "dup") is True
def test_claim_qualified_content_false_once_retired(session, tmp_path):
f = _make_file(tmp_path, "revalidate_retired.png")
content = create_content(session, path=f, hash="dup")
session.commit()
session.execute(
update(AssetContent).where(AssetContent.id == content.id).values(is_missing=True)
)
session.commit()
assert claim_qualified_content(session, content.id, "dup") is False, (
"a row retired after the lookup must not be reused"
)
def test_claim_qualified_content_false_once_hash_changed(session, tmp_path):
f = _make_file(tmp_path, "revalidate_rehashed.png")
content = create_content(session, path=f, hash="dup")
session.commit()
session.execute(
update(AssetContent).where(AssetContent.id == content.id).values(hash="different")
)
session.commit()
assert claim_qualified_content(session, content.id, "dup") is False, (
"a row whose recorded content changed identity must not be reused"
" for the hash that used to describe it"
)
def test_refresh_qualified_content_none_when_file_vanishes(session, tmp_path):
f = _make_file(tmp_path, "revalidate_gone.png")
content = create_content(session, path=f, hash="dup")
session.commit()
assert claim_qualified_content(session, content.id, "dup") is True
os.unlink(f)
assert refresh_qualified_content(session, content.id) is None
def test_size_mismatch_disqualifies_a_row_that_never_recorded_an_mtime(session, tmp_path):
f = _make_file(tmp_path, "size_mismatch.png", content=b"the real bytes on disk")
content = create_content(session, path=f, hash="dup", size_bytes=999_999)
session.commit()
assert content.mtime_ns is None, "fixture must exercise the unrecorded-mtime case"
assert lookup_for_view(session, "dup") is None, (
"a row whose recorded size disagrees with the file must not be served; an "
"unrecorded mtime is no reason to skip the size check too"
)

View File

@@ -1,231 +1,15 @@
"""Tests for metadata filtering logic in asset_reference queries."""
import pytest
from sqlalchemy.orm import Session
from app.assets.database.models import Asset, AssetReference, AssetReferenceMeta
from app.assets.database.queries import list_references_page
from app.assets.database.queries.asset_reference import convert_metadata_to_rows
from app.assets.helpers import get_utc_now
from app.assets.database.queries import create_content, create_record
def _make_asset(session: Session, hash_val: str) -> Asset:
asset = Asset(hash=hash_val, size_bytes=1024)
session.add(asset)
session.flush()
return asset
def test_interpretation_metadata_is_record_specific(session: Session) -> None:
content = create_content(session, "/output/shared-bytes.bin", hash="digest")
image = create_record(session, content.id, "image", mime_type="image/png")
text = create_record(session, content.id, "text", mime_type="text/plain")
image.system_metadata = {"width": 8}
text.system_metadata = {"encoding": "utf-8"}
def _make_reference(
session: Session,
asset: Asset,
name: str,
metadata: dict | None = None,
system_metadata: dict | None = None,
) -> AssetReference:
now = get_utc_now()
ref = AssetReference(
owner_id="",
name=name,
asset_id=asset.id,
user_metadata=metadata,
system_metadata=system_metadata,
created_at=now,
updated_at=now,
last_access_time=now,
)
session.add(ref)
session.flush()
# Build merged projection: {**system_metadata, **user_metadata}
merged = {**(system_metadata or {}), **(metadata or {})}
if merged:
for key, val in merged.items():
for row in convert_metadata_to_rows(key, val):
meta_row = AssetReferenceMeta(
asset_reference_id=ref.id,
key=row["key"],
ordinal=row.get("ordinal", 0),
val_str=row.get("val_str"),
val_num=row.get("val_num"),
val_bool=row.get("val_bool"),
val_json=row.get("val_json"),
)
session.add(meta_row)
session.flush()
return ref
class TestMetadataFilterByType:
"""Table-driven tests for metadata filtering by different value types."""
@pytest.mark.parametrize(
"match_meta,nomatch_meta,filter_key,filter_val",
[
# String matching
({"category": "models"}, {"category": "images"}, "category", "models"),
# Integer matching
({"epoch": 5}, {"epoch": 10}, "epoch", 5),
# Float matching
({"score": 0.95}, {"score": 0.5}, "score", 0.95),
# Boolean True matching
({"enabled": True}, {"enabled": False}, "enabled", True),
# Boolean False matching
({"enabled": False}, {"enabled": True}, "enabled", False),
],
ids=["string", "int", "float", "bool_true", "bool_false"],
)
def test_filter_matches_correct_value(
self, session: Session, match_meta, nomatch_meta, filter_key, filter_val
):
asset = _make_asset(session, "hash1")
_make_reference(session, asset, "match", match_meta)
_make_reference(session, asset, "nomatch", nomatch_meta)
session.commit()
refs, _, total = list_references_page(
session, metadata_filter={filter_key: filter_val}
)
assert total == 1
assert refs[0].name == "match"
@pytest.mark.parametrize(
"stored_meta,filter_key,filter_val",
[
# String no match
({"category": "models"}, "category", "other"),
# Int no match
({"epoch": 5}, "epoch", 99),
# Float no match
({"score": 0.5}, "score", 0.99),
],
ids=["string_no_match", "int_no_match", "float_no_match"],
)
def test_filter_returns_empty_when_no_match(
self, session: Session, stored_meta, filter_key, filter_val
):
asset = _make_asset(session, "hash1")
_make_reference(session, asset, "item", stored_meta)
session.commit()
refs, _, total = list_references_page(
session, metadata_filter={filter_key: filter_val}
)
assert total == 0
class TestMetadataFilterNull:
"""Tests for null/missing key filtering."""
@pytest.mark.parametrize(
"match_name,match_meta,nomatch_name,nomatch_meta,filter_key",
[
# Null matches missing key
("missing_key", {}, "has_key", {"optional": "value"}, "optional"),
# Null matches explicit null
("explicit_null", {"nullable": None}, "has_value", {"nullable": "present"}, "nullable"),
],
ids=["missing_key", "explicit_null"],
)
def test_null_filter_matches(
self, session: Session, match_name, match_meta, nomatch_name, nomatch_meta, filter_key
):
asset = _make_asset(session, "hash1")
_make_reference(session, asset, match_name, match_meta)
_make_reference(session, asset, nomatch_name, nomatch_meta)
session.commit()
refs, _, total = list_references_page(session, metadata_filter={filter_key: None})
assert total == 1
assert refs[0].name == match_name
class TestMetadataFilterList:
"""Tests for list-based (OR) filtering."""
def test_filter_by_list_matches_any(self, session: Session):
"""List values should match ANY of the values (OR)."""
asset = _make_asset(session, "hash1")
_make_reference(session, asset, "cat_a", {"category": "a"})
_make_reference(session, asset, "cat_b", {"category": "b"})
_make_reference(session, asset, "cat_c", {"category": "c"})
session.commit()
refs, _, total = list_references_page(session, metadata_filter={"category": ["a", "b"]})
assert total == 2
names = {r.name for r in refs}
assert names == {"cat_a", "cat_b"}
class TestMetadataFilterMultipleKeys:
"""Tests for multiple filter keys (AND semantics)."""
def test_multiple_keys_must_all_match(self, session: Session):
"""Multiple keys should ALL match (AND)."""
asset = _make_asset(session, "hash1")
_make_reference(session, asset, "match", {"type": "model", "version": 2})
_make_reference(session, asset, "wrong_type", {"type": "config", "version": 2})
_make_reference(session, asset, "wrong_version", {"type": "model", "version": 1})
session.commit()
refs, _, total = list_references_page(
session, metadata_filter={"type": "model", "version": 2}
)
assert total == 1
assert refs[0].name == "match"
class TestMetadataFilterEmptyDict:
"""Tests for empty filter behavior."""
def test_empty_filter_returns_all(self, session: Session):
asset = _make_asset(session, "hash1")
_make_reference(session, asset, "a", {"key": "val"})
_make_reference(session, asset, "b", {})
session.commit()
refs, _, total = list_references_page(session, metadata_filter={})
assert total == 2
class TestSystemMetadataProjection:
"""Tests for system_metadata merging into the filter projection."""
def test_system_metadata_keys_are_filterable(self, session: Session):
"""system_metadata keys should appear in the merged projection."""
asset = _make_asset(session, "hash1")
_make_reference(
session, asset, "with_sys",
system_metadata={"source": "scanner"},
)
_make_reference(session, asset, "without_sys")
session.commit()
refs, _, total = list_references_page(
session, metadata_filter={"source": "scanner"}
)
assert total == 1
assert refs[0].name == "with_sys"
def test_user_metadata_overrides_system_metadata(self, session: Session):
"""user_metadata should win when both have the same key."""
asset = _make_asset(session, "hash1")
_make_reference(
session, asset, "overridden",
metadata={"origin": "user_upload"},
system_metadata={"origin": "auto_scan"},
)
session.commit()
# Should match the user value, not the system value
refs, _, total = list_references_page(
session, metadata_filter={"origin": "user_upload"}
)
assert total == 1
assert refs[0].name == "overridden"
# Should NOT match the system value (it was overridden)
refs, _, total = list_references_page(
session, metadata_filter={"origin": "auto_scan"}
)
assert total == 0
assert image.content_id == text.content_id
assert image.system_metadata == {"width": 8}
assert text.system_metadata == {"encoding": "utf-8"}

View File

@@ -0,0 +1,137 @@
import os
import pytest
from sqlalchemy import create_engine, select, update
from sqlalchemy.exc import IntegrityError
from sqlalchemy.orm import Session
from app.assets.database.models import Asset, AssetContent, AssetTag
from app.assets.database.queries.records import (
RecordPageSpec,
create_content,
create_record,
delete_record,
get_record_by_id,
list_records_page,
mark_content_missing,
)
from app.database.models import Base
@pytest.fixture
def session():
engine = create_engine("sqlite:///:memory:", connect_args={"check_same_thread": False})
Base.metadata.create_all(engine)
with Session(engine) as sess:
yield sess
def test_listing_includes_missing_by_default_with_missing_tag(session):
content = create_content(session, path="/tmp/f1")
record = create_record(session, content_id=content.id, name="test")
mark_content_missing(session, content.id)
session.commit()
results, tag_map, total = list_records_page(session, RecordPageSpec())
assert any(r.id == record.id for r in results)
assert total == 1
assert "missing" in tag_map.get(record.id, [])
def test_listing_excludes_missing_with_filter(session):
content = create_content(session, path="/tmp/f2")
record = create_record(session, content_id=content.id, name="test2")
mark_content_missing(session, content.id)
session.commit()
results, _, total = list_records_page(
session,
RecordPageSpec(none_tags=("missing",)),
)
assert not any(r.id == record.id for r in results), "Missing record should be excluded"
assert total == 0
def test_delete_record_never_deletes_referenced_preview(session):
preview_content = create_content(session, path="/tmp/preview")
preview_record = create_record(session, content_id=preview_content.id, name="preview")
content1 = create_content(session, path="/tmp/f3")
content2 = create_content(session, path="/tmp/f4")
r1 = create_record(session, content_id=content1.id, name="r1")
r2 = create_record(session, content_id=content2.id, name="r2")
session.execute(update(Asset).where(Asset.id == r1.id).values(preview_id=preview_record.id))
session.execute(update(Asset).where(Asset.id == r2.id).values(preview_id=preview_record.id))
session.commit()
delete_record(session, r1.id)
session.commit()
assert get_record_by_id(session, preview_record.id) is not None, "Preview must survive"
delete_record(session, r2.id)
session.commit()
assert get_record_by_id(session, preview_record.id) is not None, (
"Preview must survive even with no referrers"
)
content_row = session.execute(
select(AssetContent).where(AssetContent.id == preview_content.id)
).scalar_one_or_none()
assert content_row is not None, "Preview content row should survive (D-3 floor)"
def test_concurrent_create_content_same_path(tmp_path):
db_path = str(tmp_path / "concurrent.db")
engine = create_engine(f"sqlite:///{db_path}", connect_args={"check_same_thread": False})
Base.metadata.create_all(engine)
shared_path = os.path.abspath("/tmp/shared")
with Session(engine) as s1, Session(engine) as s2:
create_content(s1, path=shared_path)
s1.commit()
create_content(s2, path=shared_path)
s2.commit()
with Session(engine) as s:
live_rows = list(
s.execute(
select(AssetContent).where(
AssetContent.path == shared_path,
AssetContent.is_missing.is_(False),
)
).scalars()
)
assert len(live_rows) == 1, f"Expected exactly one live row, got {len(live_rows)}"
def test_create_content_check_violation_surfaces_the_constraint_error(session):
with pytest.raises(IntegrityError):
create_content(session, path="/tmp/negative-size", size_bytes=-1)
def test_create_content_check_violation_on_mtime_surfaces(session):
with pytest.raises(IntegrityError):
create_content(session, path="/tmp/negative-mtime", size_bytes=0, mtime_ns=-1)
def test_create_content_uniqueness_race_returns_existing_live_row(session):
first = create_content(session, path="/tmp/race")
second = create_content(session, path="/tmp/race")
assert second.id == first.id
def test_create_record_links_a_repeated_tag_name_once(session):
content = create_content(session, path="/tmp/repeated-tag")
record = create_record(session, content_id=content.id, name="repeat", tags=["x", "x"])
session.commit()
linked = session.scalars(
select(AssetTag.tag_name).where(AssetTag.asset_id == record.id)
).all()
assert list(linked) == ["x"], (
"a tag name repeated within one call must link once, not raise on the composite PK"
)

View File

@@ -1,369 +1,23 @@
import pytest
from sqlalchemy.orm import Session
from app.assets.database.models import Asset, AssetReference, AssetReferenceTag, AssetReferenceMeta, Tag
from app.assets.database.queries import (
ensure_tags_exist,
get_reference_tags,
set_reference_tags,
add_tags_to_reference,
remove_tags_from_reference,
add_missing_tag_for_asset_id,
remove_missing_tag_for_asset_id,
list_tags_with_usage,
bulk_insert_tags_and_meta,
create_content,
create_record,
fetch_record_tags,
list_records_page,
)
from app.assets.helpers import get_utc_now
from app.assets.database.queries.records import RecordPageSpec
def _make_asset(session: Session, hash_val: str | None = None) -> Asset:
asset = Asset(hash=hash_val, size_bytes=1024)
session.add(asset)
session.flush()
return asset
def test_record_tags_filter_records_without_affecting_content(session: Session) -> None:
content = create_content(session, "/models/shared.safetensors")
tagged = create_record(session, content.id, "tagged", tags=["model"])
untagged = create_record(session, content.id, "untagged")
def _make_reference(session: Session, asset: Asset, name: str = "test", owner_id: str = "") -> AssetReference:
now = get_utc_now()
ref = AssetReference(
owner_id=owner_id,
name=name,
asset_id=asset.id,
created_at=now,
updated_at=now,
last_access_time=now,
records, _, _ = list_records_page(
session,
RecordPageSpec(all_tags=("model",)),
)
session.add(ref)
session.flush()
return ref
class TestEnsureTagsExist:
def test_creates_new_tags(self, session: Session):
ensure_tags_exist(session, ["alpha", "beta"])
session.commit()
tags = session.query(Tag).all()
assert {t.name for t in tags} == {"alpha", "beta"}
def test_is_idempotent(self, session: Session):
ensure_tags_exist(session, ["alpha"])
ensure_tags_exist(session, ["alpha"])
session.commit()
assert session.query(Tag).count() == 1
def test_normalizes_tags(self, session: Session):
ensure_tags_exist(session, [" ALPHA ", "Beta", "alpha"])
session.commit()
tags = session.query(Tag).all()
assert {t.name for t in tags} == {"ALPHA", "Beta", "alpha"}
def test_empty_list_is_noop(self, session: Session):
ensure_tags_exist(session, [])
session.commit()
assert session.query(Tag).count() == 0
class TestGetReferenceTags:
def test_returns_empty_for_no_tags(self, session: Session):
asset = _make_asset(session, "hash1")
ref = _make_reference(session, asset)
tags = get_reference_tags(session, reference_id=ref.id)
assert tags == []
def test_returns_tags_for_reference(self, session: Session):
asset = _make_asset(session, "hash1")
ref = _make_reference(session, asset)
ensure_tags_exist(session, ["tag1", "tag2"])
session.add_all([
AssetReferenceTag(asset_reference_id=ref.id, tag_name="tag1", origin="manual", added_at=get_utc_now()),
AssetReferenceTag(asset_reference_id=ref.id, tag_name="tag2", origin="manual", added_at=get_utc_now()),
])
session.flush()
tags = get_reference_tags(session, reference_id=ref.id)
assert set(tags) == {"tag1", "tag2"}
class TestSetReferenceTags:
def test_adds_new_tags(self, session: Session):
asset = _make_asset(session, "hash1")
ref = _make_reference(session, asset)
result = set_reference_tags(session, reference_id=ref.id, tags=["a", "b"])
session.commit()
assert set(result.added) == {"a", "b"}
assert result.removed == []
assert set(result.total) == {"a", "b"}
def test_removes_old_tags(self, session: Session):
asset = _make_asset(session, "hash1")
ref = _make_reference(session, asset)
set_reference_tags(session, reference_id=ref.id, tags=["a", "b", "c"])
result = set_reference_tags(session, reference_id=ref.id, tags=["a"])
session.commit()
assert result.added == []
assert set(result.removed) == {"b", "c"}
assert result.total == ["a"]
def test_replaces_tags(self, session: Session):
asset = _make_asset(session, "hash1")
ref = _make_reference(session, asset)
set_reference_tags(session, reference_id=ref.id, tags=["a", "b"])
result = set_reference_tags(session, reference_id=ref.id, tags=["b", "c"])
session.commit()
assert result.added == ["c"]
assert result.removed == ["a"]
assert set(result.total) == {"b", "c"}
class TestAddTagsToReference:
def test_adds_tags(self, session: Session):
asset = _make_asset(session, "hash1")
ref = _make_reference(session, asset)
result = add_tags_to_reference(session, reference_id=ref.id, tags=["x", "y"])
session.commit()
assert set(result.added) == {"x", "y"}
assert result.already_present == []
def test_reports_already_present(self, session: Session):
asset = _make_asset(session, "hash1")
ref = _make_reference(session, asset)
add_tags_to_reference(session, reference_id=ref.id, tags=["x"])
result = add_tags_to_reference(session, reference_id=ref.id, tags=["x", "y"])
session.commit()
assert result.added == ["y"]
assert result.already_present == ["x"]
def test_raises_for_missing_reference(self, session: Session):
with pytest.raises(ValueError, match="not found"):
add_tags_to_reference(session, reference_id="nonexistent", tags=["x"])
class TestRemoveTagsFromReference:
def test_removes_tags(self, session: Session):
asset = _make_asset(session, "hash1")
ref = _make_reference(session, asset)
add_tags_to_reference(session, reference_id=ref.id, tags=["a", "b", "c"])
result = remove_tags_from_reference(session, reference_id=ref.id, tags=["a", "b"])
session.commit()
assert set(result.removed) == {"a", "b"}
assert result.not_present == []
assert result.total_tags == ["c"]
def test_reports_not_present(self, session: Session):
asset = _make_asset(session, "hash1")
ref = _make_reference(session, asset)
add_tags_to_reference(session, reference_id=ref.id, tags=["a"])
result = remove_tags_from_reference(session, reference_id=ref.id, tags=["a", "x"])
session.commit()
assert result.removed == ["a"]
assert result.not_present == ["x"]
def test_raises_for_missing_reference(self, session: Session):
with pytest.raises(ValueError, match="not found"):
remove_tags_from_reference(session, reference_id="nonexistent", tags=["x"])
class TestMissingTagFunctions:
def test_add_missing_tag_for_asset_id(self, session: Session):
asset = _make_asset(session, "hash1")
ref = _make_reference(session, asset)
ensure_tags_exist(session, ["missing"])
add_missing_tag_for_asset_id(session, asset_id=asset.id)
session.commit()
tags = get_reference_tags(session, reference_id=ref.id)
assert "missing" in tags
def test_add_missing_tag_is_idempotent(self, session: Session):
asset = _make_asset(session, "hash1")
ref = _make_reference(session, asset)
ensure_tags_exist(session, ["missing"])
add_missing_tag_for_asset_id(session, asset_id=asset.id)
add_missing_tag_for_asset_id(session, asset_id=asset.id)
session.commit()
links = session.query(AssetReferenceTag).filter_by(asset_reference_id=ref.id, tag_name="missing").all()
assert len(links) == 1
def test_remove_missing_tag_for_asset_id(self, session: Session):
asset = _make_asset(session, "hash1")
ref = _make_reference(session, asset)
ensure_tags_exist(session, ["missing"])
add_missing_tag_for_asset_id(session, asset_id=asset.id)
remove_missing_tag_for_asset_id(session, asset_id=asset.id)
session.commit()
tags = get_reference_tags(session, reference_id=ref.id)
assert "missing" not in tags
class TestListTagsWithUsage:
def test_returns_tags_with_counts(self, session: Session):
ensure_tags_exist(session, ["used", "unused"])
asset = _make_asset(session, "hash1")
ref = _make_reference(session, asset)
add_tags_to_reference(session, reference_id=ref.id, tags=["used"])
session.commit()
rows, total = list_tags_with_usage(session)
tag_dict = {name: count for name, count in rows}
assert tag_dict["used"] == 1
assert tag_dict["unused"] == 0
assert total == 2
def test_exclude_zero_counts(self, session: Session):
ensure_tags_exist(session, ["used", "unused"])
asset = _make_asset(session, "hash1")
ref = _make_reference(session, asset)
add_tags_to_reference(session, reference_id=ref.id, tags=["used"])
session.commit()
rows, total = list_tags_with_usage(session, include_zero=False)
tag_names = {name for name, _ in rows}
assert "used" in tag_names
assert "unused" not in tag_names
def test_prefix_filter(self, session: Session):
ensure_tags_exist(session, ["alpha", "beta", "alphabet"])
session.commit()
rows, total = list_tags_with_usage(session, prefix="alph")
tag_names = {name for name, _ in rows}
assert tag_names == {"alpha", "alphabet"}
def test_prefix_filter_is_case_sensitive(self, session: Session):
ensure_tags_exist(session, ["model_type:LLM", "model_type:llm"])
session.commit()
rows, total = list_tags_with_usage(session, prefix="model_type:L")
tag_names = {name for name, _ in rows}
assert tag_names == {"model_type:LLM"}
assert total == 1
def test_order_by_name(self, session: Session):
ensure_tags_exist(session, ["zebra", "alpha", "middle"])
session.commit()
rows, _ = list_tags_with_usage(session, order="name_asc")
names = [name for name, _ in rows]
assert names == ["alpha", "middle", "zebra"]
def test_owner_visibility(self, session: Session):
ensure_tags_exist(session, ["shared-tag", "owner-tag"])
asset = _make_asset(session, "hash1")
shared_ref = _make_reference(session, asset, name="shared", owner_id="")
owner_ref = _make_reference(session, asset, name="owned", owner_id="user1")
add_tags_to_reference(session, reference_id=shared_ref.id, tags=["shared-tag"])
add_tags_to_reference(session, reference_id=owner_ref.id, tags=["owner-tag"])
session.commit()
# Empty owner sees only shared
rows, _ = list_tags_with_usage(session, owner_id="", include_zero=False)
tag_dict = {name: count for name, count in rows}
assert tag_dict.get("shared-tag", 0) == 1
assert tag_dict.get("owner-tag", 0) == 0
# User1 sees both
rows, _ = list_tags_with_usage(session, owner_id="user1", include_zero=False)
tag_dict = {name: count for name, count in rows}
assert tag_dict.get("shared-tag", 0) == 1
assert tag_dict.get("owner-tag", 0) == 1
class TestBulkInsertTagsAndMeta:
def test_inserts_tags(self, session: Session):
asset = _make_asset(session, "hash1")
ref = _make_reference(session, asset)
ensure_tags_exist(session, ["bulk-tag1", "bulk-tag2"])
session.commit()
now = get_utc_now()
tag_rows = [
{"asset_reference_id": ref.id, "tag_name": "bulk-tag1", "origin": "manual", "added_at": now},
{"asset_reference_id": ref.id, "tag_name": "bulk-tag2", "origin": "manual", "added_at": now},
]
bulk_insert_tags_and_meta(session, tag_rows=tag_rows, meta_rows=[])
session.commit()
tags = get_reference_tags(session, reference_id=ref.id)
assert set(tags) == {"bulk-tag1", "bulk-tag2"}
def test_inserts_meta(self, session: Session):
asset = _make_asset(session, "hash1")
ref = _make_reference(session, asset)
session.commit()
meta_rows = [
{
"asset_reference_id": ref.id,
"key": "meta-key",
"ordinal": 0,
"val_str": "meta-value",
"val_num": None,
"val_bool": None,
"val_json": None,
},
]
bulk_insert_tags_and_meta(session, tag_rows=[], meta_rows=meta_rows)
session.commit()
meta = session.query(AssetReferenceMeta).filter_by(asset_reference_id=ref.id).all()
assert len(meta) == 1
assert meta[0].key == "meta-key"
assert meta[0].val_str == "meta-value"
def test_ignores_conflicts(self, session: Session):
asset = _make_asset(session, "hash1")
ref = _make_reference(session, asset)
ensure_tags_exist(session, ["existing-tag"])
add_tags_to_reference(session, reference_id=ref.id, tags=["existing-tag"])
session.commit()
now = get_utc_now()
tag_rows = [
{"asset_reference_id": ref.id, "tag_name": "existing-tag", "origin": "duplicate", "added_at": now},
]
bulk_insert_tags_and_meta(session, tag_rows=tag_rows, meta_rows=[])
session.commit()
# Should still have only one tag link
links = session.query(AssetReferenceTag).filter_by(asset_reference_id=ref.id, tag_name="existing-tag").all()
assert len(links) == 1
# Origin should be original, not overwritten
assert links[0].origin == "manual"
def test_empty_lists_is_noop(self, session: Session):
bulk_insert_tags_and_meta(session, tag_rows=[], meta_rows=[])
assert session.query(AssetReferenceTag).count() == 0
assert session.query(AssetReferenceMeta).count() == 0
assert [record.id for record in records] == [tagged.id]
assert fetch_record_tags(session, untagged.id) == []

View File

@@ -0,0 +1,198 @@
from __future__ import annotations
from collections import Counter
from collections.abc import Sequence
from sqlalchemy.orm import Session
from app.assets.database.models import Asset, Tag
from app.assets.database.queries.records import (
RecordPageSpec,
create_content,
create_record,
list_records_page,
mark_content_missing,
)
from app.assets.database.queries.tags import (
list_tag_counts_for_filtered_assets,
list_tags_with_usage,
)
def _seed_asset(
session: Session, path: str, name: str, tags: Sequence[str]
) -> Asset:
content = create_content(session, path=path)
return create_record(session, content_id=content.id, name=name, tags=list(tags))
def test_usage_counts_match_seeded_asset_tag_rows(session: Session) -> None:
_seed_asset(session, "/models/a", "a", ["models", "checkpoint"])
_seed_asset(session, "/models/b", "b", ["models"])
_seed_asset(session, "/models/c", "c", ["lora"])
session.commit()
rows, total = list_tags_with_usage(session)
assert dict(rows) == {"models": 2, "checkpoint": 1, "lora": 1}
assert total == 3
def test_prefix_filters_tags_by_leading_substring(session: Session) -> None:
_seed_asset(session, "/m/1", "1", ["model", "mask", "zebra"])
session.commit()
rows, total = list_tags_with_usage(session, prefix="ma")
assert [name for name, _ in rows] == ["mask"]
assert total == 1
def test_limit_and_offset_paginate_over_real_total(session: Session) -> None:
for index in range(5):
_seed_asset(session, f"/p/{index}", str(index), [f"tag{index}"])
session.commit()
rows, total = list_tags_with_usage(session, limit=2, offset=1, order="name_asc")
assert [name for name, _ in rows] == ["tag1", "tag2"]
assert total == 5
def test_include_zero_false_drops_unused_tags_and_shrinks_total(session: Session) -> None:
_seed_asset(session, "/z/1", "1", ["used"])
session.add(Tag(name="unused"))
session.commit()
with_zero, total_with = list_tags_with_usage(session, include_zero=True)
without_zero, total_without = list_tags_with_usage(session, include_zero=False)
assert ("unused", 0) in with_zero
assert total_with == 2
assert [name for name, _ in without_zero] == ["used"]
assert total_without == 1
def test_order_is_count_desc_then_name_asc_or_name_asc(session: Session) -> None:
_seed_asset(session, "/o/1", "1", ["gamma"])
_seed_asset(session, "/o/2", "2", ["gamma", "beta"])
_seed_asset(session, "/o/3", "3", ["gamma", "beta", "alpha"])
session.commit()
by_count, _ = list_tags_with_usage(session, order="count_desc")
by_name, _ = list_tags_with_usage(session, order="name_asc")
assert [name for name, _ in by_count] == ["gamma", "beta", "alpha"]
assert [name for name, _ in by_name] == ["alpha", "beta", "gamma"]
def test_missing_content_counts_all_its_tags_including_missing(
session: Session,
) -> None:
_seed_asset(session, "/live", "live", ["foo"])
gone = create_content(session, path="/gone")
create_record(session, content_id=gone.id, name="gone", tags=["foo"])
mark_content_missing(session, gone.id)
session.commit()
counts = dict(list_tags_with_usage(session)[0])
assert counts["foo"] == 2
assert counts["missing"] == 1
def test_histogram_counts_every_tag_on_the_filtered_assets(session: Session) -> None:
_seed_asset(session, "/h/a", "a", ["models", "checkpoint"])
_seed_asset(session, "/h/b", "b", ["models", "lora"])
_seed_asset(session, "/h/c", "c", ["other"])
session.commit()
histogram = list_tag_counts_for_filtered_assets(session, include_tags=["models"])
assert histogram == {"models": 2, "checkpoint": 1, "lora": 1}
def test_histogram_includes_missing_content_assets(session: Session) -> None:
_seed_asset(session, "/hm/live", "live", ["models"])
gone = create_content(session, path="/hm/gone")
create_record(session, content_id=gone.id, name="gone", tags=["models"])
mark_content_missing(session, gone.id)
session.commit()
histogram = list_tag_counts_for_filtered_assets(session, include_tags=["models"])
assert histogram == {"models": 2, "missing": 1}
def test_refine_and_list_share_the_same_filtered_asset_set(session: Session) -> None:
kept = _seed_asset(session, "/s/a", "checkpoint-a", ["models", "checkpoint"])
_seed_asset(session, "/s/b", "checkpoint-b", ["models", "archived"])
_seed_asset(session, "/s/c", "lora-c", ["models"])
gone = create_content(session, path="/s/d")
missing_match = create_record(
session, content_id=gone.id, name="check-d", tags=["models", "checkpoint"]
)
mark_content_missing(session, gone.id)
session.commit()
spec = RecordPageSpec(
all_tags=("models",),
none_tags=("archived",),
name_contains="check",
limit=100,
)
records, tag_map, _ = list_records_page(session, spec)
expected = Counter(
tag for record in records for tag in tag_map.get(record.id, [])
)
histogram = list_tag_counts_for_filtered_assets(
session,
include_tags=["models"],
exclude_tags=["archived"],
name_contains="check",
limit=100,
)
assert {record.id for record in records} == {kept.id, missing_match.id}
assert histogram == dict(expected)
assert histogram == {"models": 2, "checkpoint": 2, "missing": 1}
def test_histogram_any_tags_matches_the_list_union(session: Session) -> None:
red = _seed_asset(session, "/an/a", "a", ["red"])
blue = _seed_asset(session, "/an/b", "b", ["blue"])
_seed_asset(session, "/an/c", "c", ["green"])
session.commit()
records, _, _ = list_records_page(
session, RecordPageSpec(any_tags=("red", "blue"), limit=100)
)
histogram = list_tag_counts_for_filtered_assets(session, any_tags=["red", "blue"])
assert {record.id for record in records} == {red.id, blue.id}
assert histogram == {"red": 1, "blue": 1}
def test_list_refine_and_tags_agree_on_counts_including_missing(
session: Session,
) -> None:
_seed_asset(session, "/agree/live", "live", ["models", "checkpoint"])
gone = create_content(session, path="/agree/gone")
create_record(session, content_id=gone.id, name="gone", tags=["models", "lora"])
mark_content_missing(session, gone.id)
session.commit()
records, tag_map, total = list_records_page(
session, RecordPageSpec(all_tags=("models",), limit=100)
)
list_tally = dict(
Counter(tag for record in records for tag in tag_map.get(record.id, []))
)
refine = list_tag_counts_for_filtered_assets(session, include_tags=["models"])
global_usage = dict(list_tags_with_usage(session)[0])
expected = {"models": 2, "checkpoint": 1, "lora": 1, "missing": 1}
assert total == 2
assert list_tally == expected
assert refine == expected
assert global_usage == expected

View File

@@ -1,11 +1,13 @@
import tempfile
from contextlib import contextmanager
from pathlib import Path
from unittest.mock import patch
import pytest
from sqlalchemy import create_engine, event
from sqlalchemy.orm import Session
from sqlalchemy.orm import Session, Session as SASession
from app.assets import mode
from app.assets.database.models import Base
@@ -15,6 +17,16 @@ def autoclean_unit_test_assets():
yield
@pytest.fixture(autouse=True)
def initialised_hash_mode():
class _HashingOff:
enable_asset_hashing = False
mode.init(_HashingOff())
yield
mode.init(None)
@pytest.fixture
def db_engine():
"""In-memory SQLite engine for fast unit tests."""
@@ -48,9 +60,6 @@ def session(db_engine):
@pytest.fixture
def mock_create_session(db_engine):
"""Patch create_session to use our in-memory database."""
from contextlib import contextmanager
from sqlalchemy.orm import Session as SASession
@contextmanager
def _create_session():
with SASession(db_engine) as sess:

View File

@@ -0,0 +1,35 @@
from __future__ import annotations
import os
def stored_path(path: str) -> str:
return os.path.abspath(path)
def prefix_case_paths(root: str) -> list[tuple[str, bool]]:
parent = os.path.dirname(root)
name = os.path.basename(root)
upper = name.upper()
assert upper != name, f"root basename {name!r} has no distinct uppercase form"
return [
(root, True),
(os.path.join(root, "child.png"), True),
(os.path.join(root, "sub", "deep.png"), True),
(os.path.join(root, "meta%_*?[].png"), True),
(os.path.join(parent, upper), False),
(os.path.join(parent, upper, "case.png"), False),
(root + "-other" + os.sep + "lexical.png", False),
(root + "extra", False),
(os.path.join(parent, "unrelated.png"), False),
(os.path.join(root, os.pardir, "escaped.png"), False),
(os.path.join(root, "sub", os.pardir, "in.png"), True),
(os.path.join(name, "relative.png"), False),
(root + os.sep + os.sep + "doubled.png", True),
(os.path.join(root, "trailing.png") + os.sep, True),
]
def expected_prefix_case_paths(root: str) -> set[str]:
return {os.path.abspath(path) for path, expected in prefix_case_paths(root) if expected}

View File

@@ -0,0 +1,170 @@
import os
from collections.abc import Iterator
from pathlib import Path
from unittest.mock import patch
import pytest
from sqlalchemy import select
from app.assets import scanner_admission
from app.assets.database.models import AssetContent
from app.assets.scanner import (
_WATCH_LIST,
_WatchEntry,
_should_skip_extension,
_two_stat_admit,
tick_watch_list,
)
@pytest.fixture(autouse=True)
def clear_watch_list() -> Iterator[None]:
_WATCH_LIST.clear()
yield
_WATCH_LIST.clear()
def test_part_suffix_never_admitted():
assert _should_skip_extension("model.safetensors.part") is True
def test_size_drift_watch_listed(temp_dir: Path, monkeypatch):
path = temp_dir / "download.bin"
path.write_bytes(b"first")
first_stat = path.stat()
def changed_stat(_: float) -> None:
path.write_bytes(b"changed")
monkeypatch.setattr("app.assets.scanner_admission.time.sleep", changed_stat)
admitted, watched = _two_stat_admit([(str(path), first_stat)])
assert admitted == []
assert watched == [str(path)]
def test_watch_list_is_bounded_when_distinct_paths_keep_changing(temp_dir: Path, monkeypatch):
candidates: list[tuple[str, os.stat_result]] = []
for index in range(1000):
path = temp_dir / f"changing-{index}.bin"
path.write_bytes(b"before")
candidates.append((str(path), path.stat()))
path.write_bytes(b"after-change")
monkeypatch.setattr("app.assets.scanner_admission.time.sleep", lambda _: None)
_two_stat_admit(candidates)
assert scanner_admission._WATCH_LIST_MAX_SIZE == 256
assert len(_WATCH_LIST) <= scanner_admission._WATCH_LIST_MAX_SIZE
def test_refreshing_watched_path_preserves_ticks_and_replaces_stat(temp_dir: Path, monkeypatch):
path = temp_dir / "changing.bin"
path.write_bytes(b"first")
first_stat = path.stat()
path.write_bytes(b"second-version")
monkeypatch.setattr("app.assets.scanner_admission.time.sleep", lambda _: None)
_two_stat_admit([(str(path), first_stat)])
_WATCH_LIST[0].ticks = 7
refresh_first_stat = path.stat()
path.write_bytes(b"third-version-is-longer")
refreshed_stat = path.stat()
_two_stat_admit([(str(path), refresh_first_stat)])
assert len(_WATCH_LIST) == 1
assert _WATCH_LIST[0].last_stat == refreshed_stat
assert _WATCH_LIST[0].ticks == 7
def test_never_stabilizes_dropped_after_cap(session, temp_dir: Path):
path = temp_dir / "moving.bin"
path.write_bytes(b"0")
initial_stat = path.stat()
previous_target_ns = initial_stat.st_mtime_ns
_WATCH_LIST[:] = [_WatchEntry(str(path), initial_stat)]
for index in range(scanner_admission._WATCH_SCAN_RETRIES):
path.write_bytes(str(index + 1).encode())
target_ns = max(path.stat().st_mtime_ns, previous_target_ns) + 1_000_000
os.utime(path, ns=(target_ns, target_ns))
previous_target_ns = target_ns
tick_watch_list(session)
assert _WATCH_LIST == []
assert session.scalars(select(AssetContent)).all() == []
def test_stable_scan_admission_removes_watch_entry_before_next_tick(session, temp_dir: Path, monkeypatch):
path = temp_dir / "stable.bin"
path.write_bytes(b"complete")
current_stat = path.stat()
_WATCH_LIST[:] = [_WatchEntry(str(path), current_stat, ticks=4)]
monkeypatch.setattr("app.assets.scanner_admission.time.sleep", lambda _: None)
admitted, watched = _two_stat_admit([(str(path), current_stat)])
entries_after_admission = len(_WATCH_LIST)
with (
patch("app.assets.scanner_admission.compute_loader_path", return_value="stable.bin"),
patch(
"app.assets.scanner_admission.get_name_and_tags_from_asset_path",
return_value=("stable.bin", []),
),
patch("app.assets.scanner.seed_asset_specs") as seed_asset_specs,
):
tick_watch_list(session)
assert admitted == [str(path)]
assert watched == []
assert entries_after_admission == 0
seed_asset_specs.assert_not_called()
def test_evicted_path_is_admitted_by_later_stable_scan(temp_dir: Path, monkeypatch):
monkeypatch.setattr(scanner_admission, "_WATCH_LIST_MAX_SIZE", 2, raising=False)
monkeypatch.setattr("app.assets.scanner_admission.time.sleep", lambda _: None)
paths: list[Path] = []
candidates: list[tuple[str, os.stat_result]] = []
for index in range(3):
path = temp_dir / f"overflow-{index}.bin"
path.write_bytes(b"before")
candidates.append((str(path), path.stat()))
path.write_bytes(b"after-change")
paths.append(path)
_two_stat_admit(candidates)
entries_after_overflow = [entry.path for entry in _WATCH_LIST]
stable_stat = paths[0].stat()
admitted, watched = _two_stat_admit([(str(paths[0]), stable_stat)])
assert entries_after_overflow == [str(paths[1]), str(paths[2])]
assert admitted == [str(paths[0])]
assert watched == []
def test_empty_candidate_batch_returns_without_paying_stability_gap(monkeypatch):
sleeps: list[float] = []
monkeypatch.setattr(
"app.assets.scanner_admission.time.sleep", lambda seconds: sleeps.append(seconds)
)
admitted, watched = _two_stat_admit([])
assert admitted == []
assert watched == []
assert sleeps == []
def test_nonempty_candidate_batch_still_pays_stability_gap(temp_dir: Path, monkeypatch):
path = temp_dir / "stable.bin"
path.write_bytes(b"complete")
first_stat = path.stat()
sleeps: list[float] = []
monkeypatch.setattr(
"app.assets.scanner_admission.time.sleep", lambda seconds: sleeps.append(seconds)
)
admitted, watched = _two_stat_admit([(str(path), first_stat)])
assert sleeps == [0.1]
assert admitted == [str(path)]
assert watched == []

View File

@@ -0,0 +1,20 @@
from unittest.mock import patch
from sqlalchemy import select
from app.assets.database.models import AssetContent
from app.assets.scanner import build_asset_specs, seed_asset_specs
def test_drifting_file_never_reaches_seed(session, temp_dir):
path = temp_dir / "drifting.bin"
path.write_bytes(b"still downloading")
with (
patch("folder_paths.get_input_directory", return_value=str(temp_dir)),
patch("app.assets.scanner._two_stat_admit", return_value=([], [str(path)])),
):
specs, _, _ = build_asset_specs([str(path)], set(), enable_metadata_extraction=False)
seed_asset_specs(session, specs)
assert list(session.scalars(select(AssetContent))) == []

View File

@@ -0,0 +1,155 @@
import json
from collections.abc import Callable
from unittest.mock import AsyncMock
import pytest
from aiohttp import web
from aiohttp.test_utils import make_mocked_request
from sqlalchemy.orm import Session
from app.assets import mode
from app.assets.api import routes
from app.assets.database.queries.records import (
create_content,
create_record,
mark_content_missing,
)
def _session_factory(engine) -> Callable[[], Session]:
return lambda: Session(engine)
@pytest.mark.asyncio
async def test_unfiltered_listing_includes_missing_entity(
db_engine, session, temp_dir, monkeypatch
):
content = create_content(session, path=str(temp_dir / "missing.png"))
record = create_record(
session,
content_id=content.id,
name="missing.png",
mime_type="image/png",
tags=["input"],
)
mark_content_missing(session, content.id)
session.commit()
monkeypatch.setattr(routes, "create_session", _session_factory(db_engine))
monkeypatch.setattr(routes, "_ASSETS_ENABLED", True)
response = await routes.list_assets_route(
make_mocked_request("GET", "/api/assets")
)
assert isinstance(response, web.Response)
response_body = response.body
assert isinstance(response_body, bytes | bytearray)
body = json.loads(response_body)
assert record.id in {item["id"] for item in body["assets"]}
assert body["total"] == 1
listed = next(item for item in body["assets"] if item["id"] == record.id)
assert "missing" in listed["tags"]
@pytest.mark.asyncio
async def test_exclude_tags_missing_hides_it(db_engine, session, temp_dir, monkeypatch):
content = create_content(session, path=str(temp_dir / "missing.png"))
record = create_record(
session,
content_id=content.id,
name="missing.png",
mime_type="image/png",
)
mark_content_missing(session, content.id)
session.commit()
monkeypatch.setattr(routes, "create_session", _session_factory(db_engine))
monkeypatch.setattr(routes, "_ASSETS_ENABLED", True)
response = await routes.list_assets_route(
make_mocked_request("GET", "/api/assets?exclude_tags=missing")
)
assert isinstance(response, web.Response)
response_body = response.body
assert isinstance(response_body, bytes | bytearray)
body = json.loads(response_body)
assert record.id not in {item["id"] for item in body["assets"]}
@pytest.mark.asyncio
async def test_listing_serves_is_immutable_as_false(
db_engine, session, temp_dir, monkeypatch
):
content = create_content(session, path=str(temp_dir / "mutable.png"))
record = create_record(
session,
content_id=content.id,
name="mutable.png",
mime_type="image/png",
)
session.commit()
monkeypatch.setattr(routes, "create_session", _session_factory(db_engine))
monkeypatch.setattr(routes, "_ASSETS_ENABLED", True)
response = await routes.list_assets_route(
make_mocked_request("GET", "/api/assets")
)
assert isinstance(response, web.Response)
response_body = response.body
assert isinstance(response_body, bytes | bytearray)
listed = next(
item
for item in json.loads(response_body)["assets"]
if item["id"] == record.id
)
assert listed["is_immutable"] is False, (
"the field must be present in the serialized body, not dropped as a default"
)
@pytest.mark.asyncio
async def test_detail_serves_is_immutable_as_false(
mock_create_session, db_engine, session, temp_dir, monkeypatch
):
content = create_content(session, path=str(temp_dir / "detail.png"))
record = create_record(
session,
content_id=content.id,
name="detail.png",
mime_type="image/png",
)
session.commit()
monkeypatch.setattr(routes, "create_session", _session_factory(db_engine))
monkeypatch.setattr(routes, "_ASSETS_ENABLED", True)
response = await routes.get_asset_route(
make_mocked_request(
"GET", f"/api/assets/{record.id}", match_info={"id": record.id}
)
)
assert isinstance(response, web.Response)
response_body = response.body
assert isinstance(response_body, bytes | bytearray)
assert json.loads(response_body)["is_immutable"] is False
@pytest.mark.asyncio
async def test_from_hash_off_mode_returns_400(monkeypatch):
monkeypatch.setattr(mode, "hashing_enabled", lambda: False)
monkeypatch.setattr(routes, "_ASSETS_ENABLED", True)
request = AsyncMock(spec=web.Request)
request.json.return_value = {"hash": f"blake3:{'a' * 64}"}
response = await routes.create_asset_from_hash_route(request)
assert isinstance(response, web.Response)
assert response.status == 400
response_body = response.body
assert isinstance(response_body, bytes | bytearray)
assert json.loads(response_body)["error"]["code"] == "FEATURE_DISABLED"

View File

@@ -1,310 +1,14 @@
"""Tests for asset_management services."""
import pytest
from sqlalchemy.orm import Session
from app.assets.database.models import Asset, AssetReference
from app.assets.database.queries import ensure_tags_exist, add_tags_to_reference
from app.assets.helpers import get_utc_now
from app.assets.services import (
get_asset_detail,
update_asset_metadata,
delete_asset_reference,
set_asset_preview,
)
from app.assets.services.asset_management import resolve_hash_to_path
from app.assets.database.queries import create_content, create_record, delete_record, get_record_by_id
def _make_asset(session: Session, hash_val: str = "blake3:test", size: int = 1024) -> Asset:
asset = Asset(hash=hash_val, size_bytes=size, mime_type="application/octet-stream")
session.add(asset)
session.flush()
return asset
def test_delete_removes_only_the_requested_delivery_record(session: Session) -> None:
content = create_content(session, "/output/shared.png")
first = create_record(session, content.id, "first.png")
second = create_record(session, content.id, "second.png")
delete_record(session, first.id)
def _make_reference(
session: Session,
asset: Asset,
name: str = "test",
owner_id: str = "",
) -> AssetReference:
now = get_utc_now()
ref = AssetReference(
owner_id=owner_id,
name=name,
asset_id=asset.id,
created_at=now,
updated_at=now,
last_access_time=now,
)
session.add(ref)
session.flush()
return ref
class TestGetAssetDetail:
def test_returns_none_for_nonexistent(self, mock_create_session):
result = get_asset_detail(reference_id="nonexistent")
assert result is None
def test_returns_asset_with_tags(self, mock_create_session, session: Session):
asset = _make_asset(session)
ref = _make_reference(session, asset, name="test.bin")
ensure_tags_exist(session, ["alpha", "beta"])
add_tags_to_reference(session, reference_id=ref.id, tags=["alpha", "beta"])
session.commit()
result = get_asset_detail(reference_id=ref.id)
assert result is not None
assert result.ref.id == ref.id
assert result.asset.hash == asset.hash
assert set(result.tags) == {"alpha", "beta"}
def test_respects_owner_visibility(self, mock_create_session, session: Session):
asset = _make_asset(session)
ref = _make_reference(session, asset, owner_id="user1")
session.commit()
# Wrong owner cannot see
result = get_asset_detail(reference_id=ref.id, owner_id="user2")
assert result is None
# Correct owner can see
result = get_asset_detail(reference_id=ref.id, owner_id="user1")
assert result is not None
class TestUpdateAssetMetadata:
def test_updates_name(self, mock_create_session, session: Session):
asset = _make_asset(session)
ref = _make_reference(session, asset, name="old_name.bin")
ref_id = ref.id
session.commit()
update_asset_metadata(
reference_id=ref_id,
name="new_name.bin",
)
# Verify by re-fetching from DB
session.expire_all()
updated_ref = session.get(AssetReference, ref_id)
assert updated_ref.name == "new_name.bin"
def test_updates_tags(self, mock_create_session, session: Session):
asset = _make_asset(session)
ref = _make_reference(session, asset)
ensure_tags_exist(session, ["old"])
add_tags_to_reference(session, reference_id=ref.id, tags=["old"])
session.commit()
result = update_asset_metadata(
reference_id=ref.id,
tags=["new1", "new2"],
)
assert set(result.tags) == {"new1", "new2"}
assert "old" not in result.tags
def test_updates_user_metadata(self, mock_create_session, session: Session):
asset = _make_asset(session)
ref = _make_reference(session, asset)
ref_id = ref.id
session.commit()
update_asset_metadata(
reference_id=ref_id,
user_metadata={"key": "value", "num": 42},
)
# Verify by re-fetching from DB
session.expire_all()
updated_ref = session.get(AssetReference, ref_id)
assert updated_ref.user_metadata["key"] == "value"
assert updated_ref.user_metadata["num"] == 42
def test_raises_for_nonexistent(self, mock_create_session):
with pytest.raises(ValueError, match="not found"):
update_asset_metadata(reference_id="nonexistent", name="fail")
def test_raises_for_wrong_owner(self, mock_create_session, session: Session):
asset = _make_asset(session)
ref = _make_reference(session, asset, owner_id="user1")
session.commit()
with pytest.raises(PermissionError, match="not owner"):
update_asset_metadata(
reference_id=ref.id,
name="new",
owner_id="user2",
)
class TestDeleteAssetReference:
def test_soft_deletes_reference(self, mock_create_session, session: Session):
asset = _make_asset(session)
ref = _make_reference(session, asset)
ref_id = ref.id
session.commit()
result = delete_asset_reference(
reference_id=ref_id,
owner_id="",
delete_content_if_orphan=False,
)
assert result is True
# Row still exists but is marked as soft-deleted
session.expire_all()
row = session.get(AssetReference, ref_id)
assert row is not None
assert row.deleted_at is not None
def test_returns_false_for_nonexistent(self, mock_create_session):
result = delete_asset_reference(
reference_id="nonexistent",
owner_id="",
)
assert result is False
def test_returns_false_for_wrong_owner(self, mock_create_session, session: Session):
asset = _make_asset(session)
ref = _make_reference(session, asset, owner_id="user1")
ref_id = ref.id
session.commit()
result = delete_asset_reference(
reference_id=ref_id,
owner_id="user2",
)
assert result is False
assert session.get(AssetReference, ref_id) is not None
def test_keeps_asset_if_other_references_exist(self, mock_create_session, session: Session):
asset = _make_asset(session)
ref1 = _make_reference(session, asset, name="ref1")
_make_reference(session, asset, name="ref2") # Second ref keeps asset alive
asset_id = asset.id
session.commit()
delete_asset_reference(
reference_id=ref1.id,
owner_id="",
delete_content_if_orphan=True,
)
# Asset should still exist
assert session.get(Asset, asset_id) is not None
def test_deletes_orphaned_asset(self, mock_create_session, session: Session):
asset = _make_asset(session)
ref = _make_reference(session, asset)
asset_id = asset.id
ref_id = ref.id
session.commit()
delete_asset_reference(
reference_id=ref_id,
owner_id="",
delete_content_if_orphan=True,
)
# Both ref and asset should be gone
assert session.get(AssetReference, ref_id) is None
assert session.get(Asset, asset_id) is None
class TestSetAssetPreview:
def test_sets_preview(self, mock_create_session, session: Session):
asset = _make_asset(session, hash_val="blake3:main")
preview_asset = _make_asset(session, hash_val="blake3:preview")
ref = _make_reference(session, asset)
preview_ref = _make_reference(session, preview_asset, name="preview.png")
ref_id = ref.id
preview_ref_id = preview_ref.id
session.commit()
set_asset_preview(
reference_id=ref_id,
preview_reference_id=preview_ref_id,
)
# Verify by re-fetching from DB
session.expire_all()
updated_ref = session.get(AssetReference, ref_id)
assert updated_ref.preview_id == preview_ref_id
def test_clears_preview(self, mock_create_session, session: Session):
asset = _make_asset(session)
preview_asset = _make_asset(session, hash_val="blake3:preview")
ref = _make_reference(session, asset)
preview_ref = _make_reference(session, preview_asset, name="preview.png")
ref.preview_id = preview_ref.id
ref_id = ref.id
session.commit()
set_asset_preview(
reference_id=ref_id,
preview_reference_id=None,
)
# Verify by re-fetching from DB
session.expire_all()
updated_ref = session.get(AssetReference, ref_id)
assert updated_ref.preview_id is None
def test_raises_for_nonexistent_ref(self, mock_create_session):
with pytest.raises(ValueError, match="not found"):
set_asset_preview(reference_id="nonexistent")
def test_raises_for_wrong_owner(self, mock_create_session, session: Session):
asset = _make_asset(session)
ref = _make_reference(session, asset, owner_id="user1")
session.commit()
with pytest.raises(PermissionError, match="not owner"):
set_asset_preview(
reference_id=ref.id,
preview_reference_id=None,
owner_id="user2",
)
class TestResolveHashToPath:
def test_returns_none_for_unknown_hash(self, mock_create_session):
result = resolve_hash_to_path("blake3:" + "a" * 64)
assert result is None
@pytest.mark.parametrize(
"ref_owner, query_owner, expect_found",
[
("user1", "user1", True),
("user1", "user2", False),
("", "anyone", True),
("", "", True),
],
ids=[
"owner_sees_own_ref",
"other_owner_blocked",
"ownerless_visible_to_anyone",
"ownerless_visible_to_empty",
],
)
def test_owner_visibility(
self, ref_owner, query_owner, expect_found,
mock_create_session, session: Session, temp_dir,
):
f = temp_dir / "file.bin"
f.write_bytes(b"data")
asset = _make_asset(session, hash_val="blake3:" + "b" * 64)
ref = _make_reference(session, asset, name="file.bin", owner_id=ref_owner)
ref.file_path = str(f)
session.commit()
result = resolve_hash_to_path(asset.hash, owner_id=query_owner)
if expect_found:
assert result is not None
assert result.abs_path == str(f)
else:
assert result is None
assert get_record_by_id(session, first.id) is None
assert get_record_by_id(session, second.id) == second

View File

@@ -81,3 +81,29 @@ def test_all_path_fields_null_without_file_path():
assert resp.loader_path is None
assert resp.display_name is None
def test_display_name_is_serialized_and_file_path_is_not(tmp_path: Path):
models = tmp_path / "models"
ckpt = models / "checkpoints"
ckpt.mkdir(parents=True)
f = ckpt / "flux.safetensors"
f.touch()
with patch("app.assets.services.path_utils.folder_paths") as mock_fp, patch(
"app.assets.services.path_utils.get_comfy_models_folders",
return_value=[("checkpoints", [str(ckpt)], {".safetensors"})],
):
mock_fp.get_input_directory.return_value = str(tmp_path / "in")
mock_fp.get_output_directory.return_value = str(tmp_path / "out")
mock_fp.get_temp_directory.return_value = str(tmp_path / "tmp")
mock_fp.models_dir = str(models)
result = _make_result(file_path=str(f), loader_path=None)
resp = _build_asset_response(result, {})
dumped = resp.model_dump(mode="json")
assert dumped["display_name"] == "checkpoints/flux.safetensors"
assert "file_path" not in dumped, (
"the namespace-rooted storage path must not leak into this public response"
)

View File

@@ -4,7 +4,12 @@ from unittest.mock import patch
import pytest
from app.assets.api.routes import _build_asset_response
from app.assets.api.routes import _build_asset_response, _build_record_response
from app.assets.database.queries.records import (
create_content,
create_record,
mark_content_missing,
)
from app.assets.services.schemas import AssetData, AssetDetailResult, ReferenceData
_TS = datetime(2024, 1, 1, 0, 0, 0)
@@ -30,6 +35,7 @@ def _make_result(
tags: list[str] | None = None,
user_metadata: dict | None = None,
with_asset: bool = True,
is_missing: bool = False,
) -> AssetDetailResult:
ref = ReferenceData(
id=ref_id,
@@ -43,7 +49,12 @@ def _make_result(
last_access_time=_TS,
)
asset = (
AssetData(hash="blake3:abc", size_bytes=1024, mime_type=mime_type)
AssetData(
hash="blake3:abc",
size_bytes=1024,
mime_type=mime_type,
is_missing=is_missing,
)
if with_asset
else None
)
@@ -298,3 +309,77 @@ def test_no_preview_url_without_content(sandboxed_comfy_roots: Path):
)
assert resp.preview_url is None, "no asset row means there is nothing to preview"
def test_no_self_preview_url_when_content_is_missing(sandboxed_comfy_roots: Path):
resp = _build_asset_response(
_make_result(
name="gone.png",
file_path=str(sandboxed_comfy_roots / "output" / "gone.png"),
mime_type="image/png",
is_missing=True,
),
{},
)
assert resp.preview_url is None, (
"a missing-content record must not advertise a preview of its own bytes; "
"those bytes are gone, so /api/view would 404"
)
def test_missing_content_still_shows_a_nominated_preview(sandboxed_comfy_roots: Path):
result = _make_result(
name="gone.safetensors",
file_path=str(sandboxed_comfy_roots / "models" / "checkpoints" / "gone.safetensors"),
mime_type="application/safetensors",
preview_id="preview-ref",
is_missing=True,
)
resp = _build_asset_response(
result, {"preview-ref": str(sandboxed_comfy_roots / "output" / "thumb.png")}
)
assert resp.preview_url == "/api/view?type=output&filename=thumb.png", (
"suppression targets self-content only; a nominated (live) preview still "
"stands in for a missing record"
)
def test_record_response_shows_self_preview_url_for_live_content(
sandboxed_comfy_roots: Path, session
):
content = create_content(
session, path=str(sandboxed_comfy_roots / "output" / "live.png")
)
record = create_record(
session, content_id=content.id, name="live.png", mime_type="image/png"
)
session.commit()
resp = _build_record_response(record, [], {})
assert resp.preview_url == "/api/view?type=output&filename=live.png", (
"a live record still previews its own bytes — the positive control that "
"proves the missing-case suppression is what withholds the URL"
)
def test_record_response_has_no_self_preview_url_when_content_is_missing(
sandboxed_comfy_roots: Path, session
):
content = create_content(
session, path=str(sandboxed_comfy_roots / "output" / "gone.png")
)
record = create_record(
session, content_id=content.id, name="gone.png", mime_type="image/png"
)
mark_content_missing(session, content.id)
session.commit()
resp = _build_record_response(record, ["missing"], {})
assert resp.preview_url is None, (
"the list surface must also withhold a missing record's self-preview URL"
)

View File

@@ -1,339 +1,40 @@
"""Tests for bulk ingest services."""
import logging
import os
from pathlib import Path
from unittest.mock import patch
from sqlalchemy.orm import Session
from app.assets.database.models import Asset, AssetReference
from app.assets.database.queries import get_reference_tags
from app.assets.database.queries import create_content, create_record
from app.assets.scanner import build_asset_specs
from app.assets.services.bulk_ingest import SeedAssetSpec, batch_insert_seed_assets
class TestBatchInsertSeedAssets:
def test_locked_file_logs_typed_error_without_path(self, caplog):
locked_path = "/private/user/models/locked.safetensors"
def test_locked_file_logs_typed_error_without_path(caplog):
locked_path = "/private/user/models/locked.safetensors"
with (
patch(
"app.assets.scanner.os.stat", side_effect=PermissionError(locked_path)
),
caplog.at_level(logging.WARNING),
):
specs, tags, skipped = build_asset_specs([locked_path], set())
assert specs == []
assert tags == set()
assert skipped == 0
assert caplog.messages == [
"Asset scan error: phase=discovery_stat error_type=permission_denied"
]
assert locked_path not in caplog.text
def test_populates_mime_type_for_model_files(self, session: Session, temp_dir: Path):
"""Verify mime_type is stored in the Asset table for model files."""
file_path = temp_dir / "model.safetensors"
file_path.write_bytes(b"fake safetensors content")
specs: list[SeedAssetSpec] = [
{
"abs_path": str(file_path),
"size_bytes": 24,
"mtime_ns": 1234567890000000000,
"info_name": "Test Model",
"tags": ["models"],
"fname": "model.safetensors",
"metadata": None,
"hash": None,
"mime_type": "application/safetensors",
}
]
result = batch_insert_seed_assets(session, specs=specs, owner_id="")
assert result.inserted_refs == 1
# Verify Asset has mime_type populated
assets = session.query(Asset).all()
assert len(assets) == 1
assert assets[0].mime_type == "application/safetensors"
def test_mime_type_none_when_not_provided(self, session: Session, temp_dir: Path):
"""Verify mime_type is None when not provided in spec."""
file_path = temp_dir / "unknown.bin"
file_path.write_bytes(b"binary data")
specs: list[SeedAssetSpec] = [
{
"abs_path": str(file_path),
"size_bytes": 11,
"mtime_ns": 1234567890000000000,
"info_name": "Unknown File",
"tags": [],
"fname": "unknown.bin",
"metadata": None,
"hash": None,
"mime_type": None,
}
]
result = batch_insert_seed_assets(session, specs=specs, owner_id="")
assert result.inserted_refs == 1
assets = session.query(Asset).all()
assert len(assets) == 1
assert assets[0].mime_type is None
def test_various_model_mime_types(self, session: Session, temp_dir: Path):
"""Verify various model file types get correct mime_type."""
test_cases = [
("model.safetensors", "application/safetensors"),
("model.pt", "application/pytorch"),
("model.ckpt", "application/pickle"),
("model.gguf", "application/gguf"),
]
specs: list[SeedAssetSpec] = []
for filename, mime_type in test_cases:
file_path = temp_dir / filename
file_path.write_bytes(b"content")
specs.append(
{
"abs_path": str(file_path),
"size_bytes": 7,
"mtime_ns": 1234567890000000000,
"info_name": filename,
"tags": [],
"fname": filename,
"metadata": None,
"hash": None,
"mime_type": mime_type,
}
)
result = batch_insert_seed_assets(session, specs=specs, owner_id="")
assert result.inserted_refs == len(test_cases)
for filename, expected_mime in test_cases:
ref = session.query(AssetReference).filter_by(name=filename).first()
assert ref is not None
asset = session.query(Asset).filter_by(id=ref.asset_id).first()
assert asset.mime_type == expected_mime, f"Expected {expected_mime} for {filename}, got {asset.mime_type}"
def test_duplicate_paths_merge_tags_before_insert(
self, session: Session, temp_dir: Path
with (
patch("app.assets.scanner.os.stat", side_effect=PermissionError(locked_path)),
caplog.at_level(logging.WARNING),
):
"""Overlapping model-folder registrations can emit the same path twice."""
file_path = temp_dir / "shared.safetensors"
file_path.write_bytes(b"shared model")
specs, tags, skipped = build_asset_specs([locked_path], set())
specs: list[SeedAssetSpec] = [
{
"abs_path": str(file_path),
"size_bytes": 12,
"mtime_ns": 1234567890000000000,
"info_name": "Shared Model",
"tags": ["models", "model_type:checkpoints"],
"fname": "shared.safetensors",
"metadata": None,
"hash": None,
"mime_type": "application/safetensors",
},
{
"abs_path": str(file_path),
"size_bytes": 12,
"mtime_ns": 1234567890000000000,
"info_name": "Shared Model",
"tags": ["models", "model_type:diffusion_models"],
"fname": "shared.safetensors",
"metadata": None,
"hash": None,
"mime_type": "application/safetensors",
},
]
result = batch_insert_seed_assets(session, specs=specs, owner_id="")
assert result.inserted_refs == 1
assert result.won_paths == 1
refs = session.query(AssetReference).all()
assert len(refs) == 1
assert set(get_reference_tags(session, reference_id=refs[0].id)) == {
"models",
"model_type:checkpoints",
"model_type:diffusion_models",
}
def test_duplicate_paths_are_merged_after_abspath_normalization(
self, session: Session, temp_dir: Path, monkeypatch
):
"""The scanner may emit equivalent paths with different spelling."""
file_path = temp_dir / "same-file.safetensors"
file_path.write_bytes(b"shared model")
monkeypatch.chdir(temp_dir)
relative_path = file_path.name
absolute_path = os.path.abspath(relative_path)
specs: list[SeedAssetSpec] = [
{
"abs_path": relative_path,
"size_bytes": 12,
"mtime_ns": 1234567890000000000,
"info_name": "Shared Model",
"tags": ["models", "model_type:checkpoints"],
"fname": "same-file.safetensors",
"metadata": None,
"hash": None,
"mime_type": "application/safetensors",
},
{
"abs_path": absolute_path,
"size_bytes": 12,
"mtime_ns": 1234567890000000000,
"info_name": "Shared Model",
"tags": ["models", "model_type:diffusion_models"],
"fname": "same-file.safetensors",
"metadata": None,
"hash": None,
"mime_type": "application/safetensors",
},
]
result = batch_insert_seed_assets(session, specs=specs, owner_id="")
assert result.inserted_refs == 1
assert result.won_paths == 1
refs = session.query(AssetReference).all()
assert len(refs) == 1
assert refs[0].file_path == absolute_path
# loader_path is persisted from the spec's fname (compute_loader_path).
assert refs[0].loader_path == "same-file.safetensors"
assert set(get_reference_tags(session, reference_id=refs[0].id)) == {
"models",
"model_type:checkpoints",
"model_type:diffusion_models",
}
def test_scanner_duplicate_shared_model_paths_keep_all_model_type_tags(
self, session: Session, temp_dir: Path
):
"""Shared extra model roots make scanner collection emit duplicate paths."""
shared_root = temp_dir / "shared"
input_dir = temp_dir / "input"
output_dir = temp_dir / "output"
temp_root = temp_dir / "temp"
for directory in (shared_root, input_dir, output_dir, temp_root):
directory.mkdir()
file_path = shared_root / "dual_use_model.safetensors"
file_path.write_bytes(b"shared model")
with (
patch("app.assets.services.path_utils.folder_paths") as mock_fp,
patch(
"app.assets.services.path_utils.get_comfy_models_folders",
return_value=[
("checkpoints", [str(shared_root)], {".safetensors"}),
("diffusion_models", [str(shared_root)], {".safetensors"}),
],
),
):
mock_fp.get_input_directory.return_value = str(input_dir)
mock_fp.get_output_directory.return_value = str(output_dir)
mock_fp.get_temp_directory.return_value = str(temp_root)
specs, tag_pool, skipped = build_asset_specs(
paths=[str(file_path), str(file_path)],
existing_paths=set(),
enable_metadata_extraction=False,
compute_hashes=False,
)
assert skipped == 0
assert len(specs) == 2
assert tag_pool == {
"models",
"model_type:checkpoints",
"model_type:diffusion_models",
}
result = batch_insert_seed_assets(session, specs=specs, owner_id="")
assert result.inserted_refs == 1
assert result.won_paths == 1
refs = session.query(AssetReference).all()
assert len(refs) == 1
assert set(get_reference_tags(session, reference_id=refs[0].id)) == {
"models",
"model_type:checkpoints",
"model_type:diffusion_models",
}
def test_loader_path_persisted_as_null_when_fname_is_none(
self, session: Session, temp_dir: Path
):
"""A file with no in-root loader path (fname=None, e.g. an orphan under
models_root) persists loader_path as NULL rather than a synthesized value."""
file_path = temp_dir / "orphan.bin"
file_path.write_bytes(b"x")
specs: list[SeedAssetSpec] = [
{
"abs_path": str(file_path),
"size_bytes": 1,
"mtime_ns": 1234567890000000000,
"info_name": "orphan.bin",
"tags": [],
"fname": None,
"metadata": None,
"hash": None,
"mime_type": None,
}
]
result = batch_insert_seed_assets(session, specs=specs, owner_id="")
assert result.inserted_refs == 1
refs = session.query(AssetReference).all()
assert len(refs) == 1
assert refs[0].file_path == str(file_path)
assert refs[0].loader_path is None
assert specs == []
assert tags == set()
assert skipped == 0
assert caplog.messages == [
"Asset scan error: phase=discovery_stat error_type=permission_denied"
]
assert locked_path not in caplog.text
class TestMetadataExtraction:
def test_extracts_mime_type_for_model_files(self, temp_dir: Path):
"""Verify metadata extraction returns correct mime_type for model files."""
from app.assets.services.metadata_extract import extract_file_metadata
def test_bulk_discovery_creates_a_record_for_each_path(session: Session) -> None:
discovered = [
("/models/a/model.safetensors", "a/model.safetensors"),
("/models/b/model.safetensors", "b/model.safetensors"),
]
file_path = temp_dir / "model.safetensors"
file_path.write_bytes(b"fake safetensors content")
records = [
create_record(session, create_content(session, path).id, name, loader_path=name)
for path, name in discovered
]
meta = extract_file_metadata(str(file_path))
assert meta.content_type == "application/safetensors"
def test_mime_type_for_various_model_formats(self, temp_dir: Path):
"""Verify various model file types get correct mime_type from metadata."""
from app.assets.services.metadata_extract import extract_file_metadata
test_cases = [
("model.safetensors", "application/safetensors"),
("model.sft", "application/safetensors"),
("model.pt", "application/pytorch"),
("model.pth", "application/pytorch"),
("model.ckpt", "application/pickle"),
("model.pkl", "application/pickle"),
("model.gguf", "application/gguf"),
]
for filename, expected_mime in test_cases:
file_path = temp_dir / filename
file_path.write_bytes(b"content")
meta = extract_file_metadata(str(file_path))
assert meta.content_type == expected_mime, f"Expected {expected_mime} for {filename}, got {meta.content_type}"
assert [record.loader_path for record in records] == [name for _, name in discovered]
assert records[0].content_id != records[1].content_id

View File

@@ -0,0 +1,113 @@
import os
from pathlib import Path
from sqlalchemy import event, select
import folder_paths
from app.assets.database.models import Asset, AssetContent
from app.assets.manager import AssetsEnabled
from comfy_execution.asset_enrichment import (
register_cached_outputs,
register_executed_outputs,
)
class _ArgsStub:
enable_assets = True
enable_asset_hashing = False
def _write_output_file(name: str) -> Path:
path = Path(folder_paths.get_output_directory()) / name
path.parent.mkdir(parents=True, exist_ok=True)
path.write_bytes(b"cached output")
return path
def _output_ui(name: str) -> dict:
return {"images": [{"filename": name, "subfolder": "", "type": "output"}]}
def _wrapper(name: str, node_id: str = "1") -> dict:
return {"meta": {"node_id": node_id}, "output": _output_ui(name)}
def test_cached_replay_binds_new_record_to_existing_content(mock_create_session):
path = _write_output_file("cached-execution-bind.png")
try:
manager = AssetsEnabled(_ArgsStub())
executed = register_executed_outputs(_output_ui(path.name), "original-job", manager)
original_id = executed["images"][0]["id"]
with mock_create_session() as session:
original_content_id = session.get(Asset, original_id).content_id
wrapper = _wrapper(path.name)
enriched = register_cached_outputs(wrapper, "cached-job", manager)
cached_id = enriched["output"]["images"][0]["id"]
with mock_create_session() as session:
contents = list(
session.scalars(
select(AssetContent).where(AssetContent.path == os.path.abspath(path))
)
)
records = list(
session.scalars(
select(Asset).where(Asset.content_id == original_content_id)
)
)
assert len(contents) == 1
assert len(records) == 2
assert original_id in {r.id for r in records}
assert cached_id in {r.id for r in records}
assert {r.job_id for r in records} == {"original-job", "cached-job"}
assert "id" not in wrapper["output"]["images"][0]
finally:
path.unlink(missing_ok=True)
def test_cached_replay_does_not_update_existing_content(mock_create_session, db_engine):
path = _write_output_file("cached-execution-no-update.png")
update_statements: list[str] = []
def capture_updates(_, __, statement, ___, ____, _____):
if statement.lstrip().upper().startswith("UPDATE"):
update_statements.append(statement)
try:
manager = AssetsEnabled(_ArgsStub())
executed = register_executed_outputs(_output_ui(path.name), "original-job", manager)
original_id = executed["images"][0]["id"]
with mock_create_session() as session:
original_content = session.get(
AssetContent, session.get(Asset, original_id).content_id
)
original_state = (
original_content.id,
original_content.hash,
original_content.size_bytes,
original_content.path,
original_content.mtime_ns,
original_content.is_missing,
original_content.created_at,
)
event.listen(db_engine, "before_cursor_execute", capture_updates)
wrapper = _wrapper(path.name)
register_cached_outputs(wrapper, "cached-job", manager)
with mock_create_session() as session:
content = session.get(AssetContent, original_state[0])
current_state = (
content.id,
content.hash,
content.size_bytes,
content.path,
content.mtime_ns,
content.is_missing,
content.created_at,
)
assert current_state == original_state
assert update_statements == []
finally:
event.remove(db_engine, "before_cursor_execute", capture_updates)
path.unlink(missing_ok=True)

View File

@@ -0,0 +1,75 @@
import os
from sqlalchemy import event, select
import folder_paths
from app.assets.database.models import Asset, AssetContent
from app.assets.database.queries.records import mark_content_missing
from app.assets.services.ingest import register_cached_output, register_executed_output
def test_cached_save_creates_delivery_record(mock_create_session, db_engine):
output_dir = folder_paths.get_output_directory()
os.makedirs(output_dir, exist_ok=True)
path = os.path.join(output_dir, "test_cached_save.png")
update_count = 0
def count_updates(_, __, statement, ___, ____, _____):
nonlocal update_count
if statement.lstrip().upper().startswith("UPDATE"):
update_count += 1
event.listen(db_engine, "before_cursor_execute", count_updates)
try:
with open(path, "wb") as file:
file.write(b"pixels")
original = register_executed_output(path, job_id="executed-job")
update_count = 0
cached = register_cached_output(path, job_id="delivery-job")
with mock_create_session() as session:
records = list(
session.scalars(select(Asset).where(Asset.content_id == original.content_id))
)
assert {record.id for record in records} == {original.id, cached.id}
assert {record.job_id for record in records} == {"executed-job", "delivery-job"}
assert update_count == 0
finally:
event.remove(db_engine, "before_cursor_execute", count_updates)
if os.path.exists(path):
os.unlink(path)
def test_cached_save_against_missing_content_is_nonevent(mock_create_session):
output_dir = folder_paths.get_output_directory()
os.makedirs(output_dir, exist_ok=True)
path = os.path.join(output_dir, "test_cached_save_missing.png")
try:
with open(path, "wb") as file:
file.write(b"pixels")
original = register_executed_output(path, job_id="executed-job")
with mock_create_session() as session:
mark_content_missing(session, original.content_id)
session.commit()
cached = register_cached_output(path, job_id="delivery-job")
assert cached is None
with mock_create_session() as session:
contents = list(
session.scalars(
select(AssetContent).where(AssetContent.path == os.path.abspath(path))
)
)
assert len(contents) == 1
assert contents[0].id == original.content_id
assert contents[0].is_missing is True
records = list(
session.scalars(select(Asset).where(Asset.job_id == "delivery-job"))
)
assert records == []
finally:
if os.path.exists(path):
os.unlink(path)

View File

@@ -0,0 +1,115 @@
import os
from sqlalchemy import update
import folder_paths
from app.assets.database.models import Asset, AssetContent
from app.assets.database.queries.records import (
create_content,
create_record,
delete_record,
get_record_by_id,
mark_content_missing,
)
from app.assets.services.asset_management import delete_asset_reference
from app.assets.services.ingest import register_executed_output
def test_hard_delete_record_content_and_file_remain(mock_create_session, session, temp_dir):
content = create_content(session, path=str(temp_dir / "file.bin"), size_bytes=4)
record = create_record(session, content_id=content.id, name="file.bin")
temp_dir.joinpath("file.bin").write_bytes(b"data")
session.commit()
record_id = record.id
content_id = content.id
path = content.path
assert delete_asset_reference(record_id) is True
session.expire_all()
assert get_record_by_id(session, record_id) is None
assert session.get(AssetContent, content_id) is not None
assert os.path.isfile(path)
def test_two_records_one_content_other_unaffected(mock_create_session, session):
content = create_content(session, path="/tmp/shared.bin")
r1 = create_record(session, content_id=content.id, name="a")
r2 = create_record(session, content_id=content.id, name="b")
session.commit()
r1_id = r1.id
r2_id = r2.id
assert delete_asset_reference(r1_id) is True
session.expire_all()
assert get_record_by_id(session, r1_id) is None
assert get_record_by_id(session, r2_id) is not None
assert session.get(AssetContent, content.id) is not None
def test_delete_record_never_deletes_pointed_at_preview(mock_create_session, session):
b_content = create_content(session, path="/tmp/preview_target.png")
b = create_record(session, content_id=b_content.id, name="B")
a_content = create_content(session, path="/tmp/owner.png")
a = create_record(session, content_id=a_content.id, name="A")
session.execute(update(Asset).where(Asset.id == a.id).values(preview_id=b.id))
session.commit()
a_id = a.id
b_id = b.id
assert delete_asset_reference(a_id) is True
session.expire_all()
assert get_record_by_id(session, a_id) is None
assert get_record_by_id(session, b_id) is not None
def test_missing_content_record_deletable(mock_create_session, session):
content = create_content(session, path="/tmp/missing.bin")
record = create_record(session, content_id=content.id, name="missing.bin")
mark_content_missing(session, content.id)
session.commit()
record_id = record.id
assert delete_asset_reference(record_id) is True
session.expire_all()
assert get_record_by_id(session, record_id) is None
def test_reregister_same_path_fresh_ids(mock_create_session, monkeypatch):
output_dir = folder_paths.get_output_directory()
os.makedirs(output_dir, exist_ok=True)
f = os.path.join(output_dir, "test_delete_reregister.png")
try:
with open(f, "wb") as fh:
fh.write(b"v1")
first = register_executed_output(f, job_id="job1")
first_id = first.id
assert delete_asset_reference(first_id) is True
with open(f, "wb") as fh:
fh.write(b"v2")
second = register_executed_output(f, job_id="job2")
assert second.id != first_id
finally:
if os.path.exists(f):
os.unlink(f)
def test_delete_nonexistent_returns_false(mock_create_session):
assert delete_asset_reference("00000000-0000-0000-0000-000000000000") is False
def test_delete_order_record_before_content(mock_create_session, session):
content = create_content(session, path="/tmp/order.bin")
record = create_record(session, content_id=content.id, name="order.bin")
session.commit()
content_id = content.id
delete_record(session, record.id)
session.flush()
assert session.get(Asset, record.id) is None
assert session.get(AssetContent, content_id) is not None

View File

@@ -0,0 +1,197 @@
import os
from pathlib import Path
from unittest.mock import patch
import pytest
from sqlalchemy import select
from app.assets.database.models import Asset, AssetContent
from app.assets.helpers import to_stored_hash
from app.assets.scanner import (
clear_pending_verifications,
drain_pending_verifications,
sync_prefixes_with_filesystem,
)
from app.assets.scanner_changes import queue_pending_verification
from app.assets.services.snapshot_hash import snapshot_hash
@pytest.fixture(autouse=True)
def _clear_pending_verifications():
clear_pending_verifications()
yield
clear_pending_verifications()
def _seed_content(session, path: Path, hash_value: str | None) -> tuple[AssetContent, Asset]:
stat = path.stat()
content = AssetContent(
path=str(path),
hash=hash_value,
size_bytes=stat.st_size,
mtime_ns=stat.st_mtime_ns,
)
session.add(content)
session.flush()
record = Asset(content_id=content.id, name=path.name)
session.add(record)
session.commit()
return content, record
def _bump_mtime(path: Path) -> None:
stat = path.stat()
os.utime(path, ns=(stat.st_atime_ns, stat.st_mtime_ns + 1_000_000))
def _stored_hash(path: Path) -> str:
snapshot = snapshot_hash(str(path))
assert snapshot is not None
digest, _ = snapshot
return to_stored_hash(digest)
def test_off_mode_same_size_touch_does_not_split(session, temp_dir: Path):
input_root = temp_dir / "input"
input_root.mkdir()
path = input_root / "touched.bin"
path.write_bytes(b"same bytes")
old_content, _ = _seed_content(session, path, hash_value="historical")
_bump_mtime(path)
with (
patch("folder_paths.get_input_directory", return_value=str(input_root)),
patch("app.assets.scanner.mode.hashing_enabled", return_value=False),
):
sync_prefixes_with_filesystem(session, [str(input_root)])
session.commit()
contents = list(session.scalars(select(AssetContent)))
assert len(contents) == 1
assert session.get(AssetContent, old_content.id).is_missing is False
def test_off_mode_size_change_splits(session, temp_dir: Path):
input_root = temp_dir / "input"
input_root.mkdir()
path = input_root / "grown.bin"
path.write_bytes(b"small")
previous_target_ns = path.stat().st_mtime_ns
old_content, _ = _seed_content(session, path, hash_value="historical")
path.write_bytes(b"a decidedly larger set of bytes")
target_ns = max(path.stat().st_mtime_ns, previous_target_ns) + 1_000_000
os.utime(path, ns=(target_ns, target_ns))
with (
patch("folder_paths.get_input_directory", return_value=str(input_root)),
patch("app.assets.scanner.mode.hashing_enabled", return_value=False),
):
sync_prefixes_with_filesystem(session, [str(input_root)])
session.commit()
contents = list(session.scalars(select(AssetContent).order_by(AssetContent.created_at)))
assert len(contents) == 2
assert session.get(AssetContent, old_content.id).is_missing is True
assert [content.is_missing for content in contents] == [True, False]
def test_hash_mode_touch_refreshes_mtime(session, temp_dir: Path):
input_root = temp_dir / "input"
input_root.mkdir()
path = input_root / "touched.bin"
path.write_bytes(b"same bytes")
old_content, _ = _seed_content(session, path, _stored_hash(path))
_bump_mtime(path)
with (
patch("folder_paths.get_input_directory", return_value=str(input_root)),
patch("app.assets.scanner.mode.hashing_enabled", return_value=True),
):
sync_prefixes_with_filesystem(session, [str(input_root)])
processed = drain_pending_verifications(session)
session.commit()
refreshed = session.get(AssetContent, old_content.id)
assert processed == 1
assert refreshed.is_missing is False
assert refreshed.mtime_ns == path.stat().st_mtime_ns
assert len(session.scalars(select(AssetContent)).all()) == 1
def test_hash_mode_real_edit_splits(session, temp_dir: Path):
input_root = temp_dir / "input"
input_root.mkdir()
path = input_root / "edited.bin"
path.write_bytes(b"old bytes")
previous_target_ns = path.stat().st_mtime_ns
old_content, _ = _seed_content(session, path, _stored_hash(path))
path.write_bytes(b"new bytes with a different length")
target_ns = max(path.stat().st_mtime_ns, previous_target_ns) + 1_000_000
os.utime(path, ns=(target_ns, target_ns))
with (
patch("folder_paths.get_input_directory", return_value=str(input_root)),
patch("app.assets.scanner.mode.hashing_enabled", return_value=True),
):
sync_prefixes_with_filesystem(session, [str(input_root)])
drain_pending_verifications(session)
session.commit()
contents = list(session.scalars(select(AssetContent).order_by(AssetContent.created_at)))
assert len(contents) == 2
assert session.get(AssetContent, old_content.id).is_missing is True
assert next(content for content in contents if not content.is_missing).hash == _stored_hash(path)
def test_old_record_id_resolves_to_missing_content_after_split(session, temp_dir: Path):
input_root = temp_dir / "input"
input_root.mkdir()
path = input_root / "edited.bin"
path.write_bytes(b"old bytes")
previous_target_ns = path.stat().st_mtime_ns
old_content, old_record = _seed_content(session, path, _stored_hash(path))
path.write_bytes(b"replacement bytes")
target_ns = max(path.stat().st_mtime_ns, previous_target_ns) + 1_000_000
os.utime(path, ns=(target_ns, target_ns))
with (
patch("folder_paths.get_input_directory", return_value=str(input_root)),
patch("app.assets.scanner.mode.hashing_enabled", return_value=True),
):
sync_prefixes_with_filesystem(session, [str(input_root)])
drain_pending_verifications(session)
session.commit()
session.expire_all()
original_record = session.get(Asset, old_record.id)
assert original_record is not None
assert original_record.content_id == old_content.id
assert original_record.content.is_missing is True
def test_hash_mode_split_uses_stat_from_the_verified_snapshot(session, temp_dir: Path, monkeypatch):
input_root = temp_dir / "input"
input_root.mkdir()
path = input_root / "changed.bin"
monkeypatch.setattr("folder_paths.get_input_directory", lambda: str(input_root))
path.write_bytes(b"old")
content, _ = _seed_content(session, path, _stored_hash(path))
queue_pending_verification(content.id)
new_payload = b"new bytes with a different size"
real_snapshot_hash = snapshot_hash
def mutate_then_hash(candidate_path: str):
path.write_bytes(new_payload)
return real_snapshot_hash(candidate_path)
monkeypatch.setattr("app.assets.scanner_changes.snapshot_hash", mutate_then_hash)
processed = drain_pending_verifications(session)
live_content = session.scalar(
select(AssetContent).where(AssetContent.is_missing.is_(False))
)
assert processed == 1
assert live_content is not None
assert live_content.size_bytes == len(new_payload)
assert live_content.mtime_ns == path.stat().st_mtime_ns

View File

@@ -1,216 +1,18 @@
"""Tests for asset enrichment (mime_type and hash population)."""
import os
from pathlib import Path
from sqlalchemy.orm import Session
from app.assets.database.models import Asset, AssetReference
from app.assets.services.file_utils import get_mtime_ns
from app.assets.scanner import (
ENRICHMENT_HASHED,
ENRICHMENT_METADATA,
ENRICHMENT_STUB,
enrich_asset,
)
from app.assets.database.queries import create_content, create_record
def _create_stub_asset(
session: Session,
file_path: str,
asset_id: str = "test-asset-id",
reference_id: str = "test-ref-id",
name: str | None = None,
) -> tuple[Asset, AssetReference]:
"""Create a stub asset with reference for testing enrichment."""
# Use the real file's mtime so the optimistic guard in enrich_asset passes
try:
stat_result = os.stat(file_path, follow_symlinks=True)
mtime_ns = get_mtime_ns(stat_result)
except OSError:
mtime_ns = 1234567890000000000
def test_unenriched_content_is_represented_by_a_null_hash(session: Session) -> None:
content = create_content(session, "/models/pending.safetensors", hash=None)
record = create_record(session, content.id, "pending.safetensors")
asset = Asset(
id=asset_id,
hash=None,
size_bytes=100,
mime_type=None,
)
session.add(asset)
session.flush()
ref = AssetReference(
id=reference_id,
asset_id=asset_id,
name=name or f"test-asset-{asset_id}",
owner_id="system",
file_path=file_path,
mtime_ns=mtime_ns,
enrichment_level=ENRICHMENT_STUB,
)
session.add(ref)
session.flush()
return asset, ref
assert record.content_id == content.id
assert content.hash is None
class TestEnrichAsset:
def test_extracts_mime_type_and_updates_asset(
self, db_engine, temp_dir: Path, session: Session
):
"""Verify mime_type is written to the Asset table during enrichment."""
file_path = temp_dir / "model.safetensors"
file_path.write_bytes(b"\x00" * 100)
def test_enrichment_does_not_merge_equal_hash_content_rows(session: Session) -> None:
first = create_content(session, "/models/one.safetensors", hash="digest")
second = create_content(session, "/models/two.safetensors", hash="digest")
asset, ref = _create_stub_asset(
session, str(file_path), "asset-1", "ref-1"
)
session.commit()
new_level = enrich_asset(
session,
file_path=str(file_path),
reference_id=ref.id,
asset_id=asset.id,
extract_metadata=True,
compute_hash=False,
)
assert new_level == ENRICHMENT_METADATA
session.expire_all()
updated_asset = session.get(Asset, "asset-1")
assert updated_asset is not None
assert updated_asset.mime_type == "application/safetensors"
def test_computes_hash_and_updates_asset(
self, db_engine, temp_dir: Path, session: Session
):
"""Verify hash is written to the Asset table during enrichment."""
file_path = temp_dir / "data.bin"
file_path.write_bytes(b"test content for hashing")
asset, ref = _create_stub_asset(
session, str(file_path), "asset-2", "ref-2"
)
session.commit()
new_level = enrich_asset(
session,
file_path=str(file_path),
reference_id=ref.id,
asset_id=asset.id,
extract_metadata=True,
compute_hash=True,
)
assert new_level == ENRICHMENT_HASHED
session.expire_all()
updated_asset = session.get(Asset, "asset-2")
assert updated_asset is not None
assert updated_asset.hash is not None
assert updated_asset.hash.startswith("blake3:")
def test_enrichment_updates_both_mime_and_hash(
self, db_engine, temp_dir: Path, session: Session
):
"""Verify both mime_type and hash are set when full enrichment runs."""
file_path = temp_dir / "model.safetensors"
file_path.write_bytes(b"\x00" * 50)
asset, ref = _create_stub_asset(
session, str(file_path), "asset-3", "ref-3"
)
session.commit()
enrich_asset(
session,
file_path=str(file_path),
reference_id=ref.id,
asset_id=asset.id,
extract_metadata=True,
compute_hash=True,
)
session.expire_all()
updated_asset = session.get(Asset, "asset-3")
assert updated_asset is not None
assert updated_asset.mime_type == "application/safetensors"
assert updated_asset.hash is not None
assert updated_asset.hash.startswith("blake3:")
def test_missing_file_returns_stub_level(
self, db_engine, temp_dir: Path, session: Session
):
"""Verify missing files don't cause errors and return STUB level."""
file_path = temp_dir / "nonexistent.bin"
asset, ref = _create_stub_asset(
session, str(file_path), "asset-4", "ref-4"
)
session.commit()
new_level = enrich_asset(
session,
file_path=str(file_path),
reference_id=ref.id,
asset_id=asset.id,
extract_metadata=True,
compute_hash=True,
)
assert new_level == ENRICHMENT_STUB
session.expire_all()
updated_asset = session.get(Asset, "asset-4")
assert updated_asset.mime_type is None
assert updated_asset.hash is None
def test_duplicate_hash_merges_into_existing_asset(
self, db_engine, temp_dir: Path, session: Session
):
"""Verify duplicate files merge into existing asset instead of failing."""
file_path_1 = temp_dir / "file1.bin"
file_path_2 = temp_dir / "file2.bin"
content = b"identical content"
file_path_1.write_bytes(content)
file_path_2.write_bytes(content)
asset1, ref1 = _create_stub_asset(
session, str(file_path_1), "asset-dup-1", "ref-dup-1"
)
asset2, ref2 = _create_stub_asset(
session, str(file_path_2), "asset-dup-2", "ref-dup-2"
)
session.commit()
enrich_asset(
session,
file_path=str(file_path_1),
reference_id=ref1.id,
asset_id=asset1.id,
extract_metadata=True,
compute_hash=True,
)
enrich_asset(
session,
file_path=str(file_path_2),
reference_id=ref2.id,
asset_id=asset2.id,
extract_metadata=True,
compute_hash=True,
)
session.expire_all()
updated_asset1 = session.get(Asset, "asset-dup-1")
assert updated_asset1 is not None
assert updated_asset1.hash is not None
updated_asset2 = session.get(Asset, "asset-dup-2")
assert updated_asset2 is None
updated_ref2 = session.get(AssetReference, "ref-dup-2")
assert updated_ref2 is not None
assert updated_ref2.asset_id == "asset-dup-1"
assert first.id != second.id

View File

@@ -0,0 +1,308 @@
from __future__ import annotations
import os
from contextlib import contextmanager
from collections.abc import Iterator
from pathlib import Path
from unittest.mock import patch
import sqlalchemy as sa
from sqlalchemy.orm import Session
from app.assets.database.queries import create_content, create_record, mark_content_missing
from app.assets.scanner import get_unenriched_assets_for_roots
from app.assets.scanner_changes import (
is_path_under_prefixes,
live_contents_under_prefixes,
)
from .path_prefix_cases import expected_prefix_case_paths, prefix_case_paths
@contextmanager
def _reuse_session(session: Session) -> Iterator[Session]:
yield session
def _raw_system_metadata(session: Session, record_id: str) -> object:
return session.execute(
sa.text("SELECT system_metadata FROM assets WHERE id = :id"),
{"id": record_id},
).scalar()
def test_create_record_stores_none_system_metadata_as_sql_null(
session: Session,
) -> None:
content = create_content(session, "/models/no-meta.safetensors", hash=None)
record = create_record(session, content.id, "no-meta.safetensors")
session.commit()
assert _raw_system_metadata(session, record.id) is None
def test_off_mode_returns_metadata_less_seeded_asset(
session: Session, temp_dir: Path
) -> None:
path = temp_dir / "fresh.safetensors"
content = create_content(session, str(path), hash=None)
record = create_record(session, content.id, path.name)
session.commit()
record_id = record.id
with (
patch("app.assets.scanner.create_session", lambda: _reuse_session(session)),
patch(
"app.assets.scanner.get_scan_prefixes_for_root",
return_value=[str(temp_dir)],
),
):
rows = get_unenriched_assets_for_roots(("models",), compute_hashes=False)
assert record_id in {row.record_id for row in rows}
def test_off_mode_excludes_asset_with_system_metadata(
session: Session, temp_dir: Path
) -> None:
path = temp_dir / "enriched.safetensors"
content = create_content(session, str(path), hash=None)
record = create_record(
session,
content.id,
path.name,
system_metadata={"architecture": "flux"},
)
session.commit()
record_id = record.id
with (
patch("app.assets.scanner.create_session", lambda: _reuse_session(session)),
patch(
"app.assets.scanner.get_scan_prefixes_for_root",
return_value=[str(temp_dir)],
),
):
rows = get_unenriched_assets_for_roots(("models",), compute_hashes=False)
assert record_id not in {row.record_id for row in rows}
def test_query_pushes_limit_into_sql(session: Session, temp_dir: Path) -> None:
for index in range(5):
path = temp_dir / f"cand-{index}.safetensors"
content = create_content(session, str(path), hash=None)
create_record(session, content.id, path.name)
session.commit()
statements: list[str] = []
def _capture(_conn, _cursor, statement, _params, _context, _executemany) -> None:
statements.append(statement)
engine = session.bind
sa.event.listen(engine, "before_cursor_execute", _capture)
try:
with (
patch("app.assets.scanner.create_session", lambda: _reuse_session(session)),
patch(
"app.assets.scanner.get_scan_prefixes_for_root",
return_value=[str(temp_dir)],
),
):
rows = get_unenriched_assets_for_roots(
("models",), compute_hashes=False, limit=2
)
finally:
sa.event.remove(engine, "before_cursor_execute", _capture)
assert len(rows) == 2
content_selects = [
s for s in statements if "asset_contents" in s.lower() and s.lower().lstrip().startswith("select")
]
assert content_selects, f"expected a SELECT over asset_contents; got {statements}"
assert any("limit" in s.lower() for s in content_selects), (
f"enrichment query must push LIMIT into SQL; got: {content_selects}"
)
def test_prefix_filter_matches_is_path_under_prefixes(
session: Session, temp_dir: Path
) -> None:
prefix = str(temp_dir)
inside_path = os.path.join(prefix, "sub", "inside.safetensors")
sibling_path = prefix + "-sibling" + os.sep + "outside.safetensors"
inside = create_record(
session, create_content(session, inside_path, hash=None).id, "inside.safetensors"
)
sibling = create_record(
session, create_content(session, sibling_path, hash=None).id, "outside.safetensors"
)
session.commit()
with (
patch("app.assets.scanner.create_session", lambda: _reuse_session(session)),
patch(
"app.assets.scanner.get_scan_prefixes_for_root",
return_value=[prefix],
),
):
returned = {
row.record_id
for row in get_unenriched_assets_for_roots(("models",), compute_hashes=False)
}
assert is_path_under_prefixes(inside_path, [prefix]) is True
assert is_path_under_prefixes(sibling_path, [prefix]) is False
assert inside.id in returned
assert sibling.id not in returned
def _seed_paths(session: Session, paths: list[str]) -> dict[str, str]:
by_record: dict[str, str] = {}
for index, path in enumerate(paths):
content = create_content(session, path, hash=None)
record = create_record(session, content.id, f"case-{index}.safetensors")
by_record[record.id] = content.path
session.commit()
return by_record
def _candidate_paths(session: Session, prefix: str) -> set[str]:
with (
patch("app.assets.scanner.create_session", lambda: _reuse_session(session)),
patch(
"app.assets.scanner.get_scan_prefixes_for_root",
return_value=[prefix],
),
):
rows = get_unenriched_assets_for_roots(("models",), compute_hashes=False)
return {row.file_path for row in rows}
def test_prefix_filter_result_set_equals_python_predicate(
session: Session, temp_dir: Path
) -> None:
root = str(temp_dir / "root")
corpus = prefix_case_paths(root)
by_record = _seed_paths(session, [path for path, _ in corpus])
stored = set(by_record.values())
returned = _candidate_paths(session, root)
expected = expected_prefix_case_paths(root)
assert returned == expected
assert expected and expected != stored
def test_prefix_holding_metacharacters_matches_only_literal_children(
session: Session, temp_dir: Path
) -> None:
root = str(temp_dir / "a_b%c*d?e[f")
inside_path = os.path.join(root, "inside.safetensors")
decoy_path = os.path.join(str(temp_dir), "aXbYcZdWeQf", "decoy.safetensors")
_seed_paths(session, [inside_path, decoy_path])
returned = _candidate_paths(session, root)
assert is_path_under_prefixes(decoy_path, [root]) is False
assert returned == {inside_path}
# --- live_contents_under_prefixes direct pins ---
# These pins exercise live_contents_under_prefixes directly, not the enrichment predicate.
def test_live_contents_under_prefixes_returns_empty_for_empty_prefixes(
session: Session, temp_dir: Path
) -> None:
_ = create_content(session, str(temp_dir / "seed.safetensors"))
returned = {content.path for content in live_contents_under_prefixes(session, [])}
assert returned == set()
def test_live_contents_under_prefixes_matches_directory_and_exact_path_prefixes(
session: Session, temp_dir: Path
) -> None:
content = create_content(session, str(temp_dir / "root" / "model.safetensors"))
under_directory = {candidate.path for candidate in live_contents_under_prefixes(session, [str(temp_dir / "root")])}
exact_path = {candidate.path for candidate in live_contents_under_prefixes(session, [content.path])}
assert under_directory == {content.path}
assert exact_path == {content.path}
def test_live_contents_under_prefixes_respects_sibling_boundary(
session: Session, temp_dir: Path
) -> None:
sibling_path = str(temp_dir / "a" / "bc" / "model.safetensors")
_ = create_content(session, sibling_path)
returned = {content.path for content in live_contents_under_prefixes(session, [str(temp_dir / "a" / "b")])}
assert returned == set()
def test_live_contents_under_prefixes_is_case_sensitive(
session: Session, temp_dir: Path
) -> None:
case_different_path = str(temp_dir / "data" / "TEMP" / "model.safetensors")
_ = create_content(session, case_different_path)
returned = {content.path for content in live_contents_under_prefixes(session, [str(temp_dir / "data" / "temp")])}
assert returned == set()
def test_live_contents_under_prefixes_treats_metacharacters_literally(
session: Session, temp_dir: Path
) -> None:
prefix = temp_dir / "a_b%c*d?e[f"
literal_child = create_content(session, str(prefix / "child.safetensors"))
decoy_path = str(temp_dir / "aXbYc*d?e[f" / "keep.safetensors")
_ = create_content(session, decoy_path)
returned = {content.path for content in live_contents_under_prefixes(session, [str(prefix)])}
assert returned == {literal_child.path}
def test_live_contents_under_prefixes_matches_any_prefix(
session: Session, temp_dir: Path
) -> None:
content = create_content(session, str(temp_dir / "second" / "model.safetensors"))
prefixes = [str(temp_dir / "first"), str(temp_dir / "second")]
returned = {candidate.path for candidate in live_contents_under_prefixes(session, prefixes)}
assert returned == {content.path}
def test_live_contents_under_prefixes_excludes_missing_content(
session: Session, temp_dir: Path
) -> None:
content = create_content(session, str(temp_dir / "root" / "missing.safetensors"))
mark_content_missing(session, content.id)
returned = {candidate.path for candidate in live_contents_under_prefixes(session, [str(temp_dir / "root")])}
assert returned == set()
def test_live_contents_under_prefixes_equals_python_predicate_for_corpus(
session: Session, temp_dir: Path
) -> None:
root = str(temp_dir / "root")
corpus = prefix_case_paths(root)
for path, _ in corpus:
_ = create_content(session, path)
prefixes = [root]
returned = {content.path for content in live_contents_under_prefixes(session, prefixes)}
assert returned == expected_prefix_case_paths(root)

View File

@@ -0,0 +1,211 @@
import os
from pathlib import Path
from unittest.mock import patch
from app.assets.database.models import Asset, AssetContent
from app.assets.helpers import to_stored_hash
from app.assets.scanner import enrich_asset
def _create_unhashed_record(session, path: Path) -> tuple[AssetContent, Asset]:
content = AssetContent(
path=str(path),
hash=None,
size_bytes=path.stat().st_size,
mtime_ns=path.stat().st_mtime_ns,
)
session.add(content)
session.flush()
record = Asset(content_id=content.id, name=path.name)
session.add(record)
session.commit()
return content, record
def test_enrichment_uses_snapshot_hash_not_direct_blake3(session, temp_dir: Path):
path = temp_dir / "stable.bin"
path.write_bytes(b"stable content")
content, record = _create_unhashed_record(session, path)
with patch(
"app.assets.scanner.snapshot_hash", return_value=("snapshot-digest", path.stat())
) as mocked_snapshot_hash:
enriched = enrich_asset(
session,
file_path=str(path),
content_id=content.id,
record_id=record.id,
extract_metadata=False,
compute_hash=True,
)
assert enriched is True
assert session.get(AssetContent, content.id).hash == to_stored_hash("snapshot-digest")
mocked_snapshot_hash.assert_called_once_with(str(path))
def test_enrichment_discards_unstable_hash(session, temp_dir: Path):
path = temp_dir / "unstable.bin"
path.write_bytes(b"unstable content")
content, record = _create_unhashed_record(session, path)
with patch("app.assets.scanner.snapshot_hash", return_value=None):
enriched = enrich_asset(
session,
file_path=str(path),
content_id=content.id,
record_id=record.id,
extract_metadata=False,
compute_hash=True,
)
assert enriched is False
assert session.get(AssetContent, content.id).hash is None
def test_enrichment_discards_metadata_read_from_a_different_file_than_the_hash(
session, temp_dir: Path
):
path = temp_dir / "swapped.bin"
path.write_bytes(b"original bytes")
content, record = _create_unhashed_record(session, path)
def _replace_file_then_hash(_file_path: str):
path.write_bytes(b"replacement bytes, a different length entirely")
return "replacement-digest", path.stat()
with patch("app.assets.scanner.snapshot_hash", side_effect=_replace_file_then_hash):
enriched = enrich_asset(
session,
file_path=str(path),
content_id=content.id,
record_id=record.id,
extract_metadata=True,
compute_hash=True,
)
assert enriched is False
session.expire_all()
assert session.get(AssetContent, content.id).hash is None
assert session.get(Asset, record.id).system_metadata is None, (
"a mismatched hash observation must discard already-computed metadata too, "
"not just the hash"
)
def test_enrichment_discards_result_when_only_the_hashed_mtime_disagrees(
session, temp_dir: Path
):
path = temp_dir / "touched.bin"
path.write_bytes(b"same length bytes")
content, record = _create_unhashed_record(session, path)
later_mtime_ns = path.stat().st_mtime_ns + 5_000_000_000
def _touch_file_then_hash(_file_path: str):
os.utime(path, ns=(later_mtime_ns, later_mtime_ns))
return "rewritten-digest", path.stat()
with patch("app.assets.scanner.snapshot_hash", side_effect=_touch_file_then_hash):
enriched = enrich_asset(
session,
file_path=str(path),
content_id=content.id,
record_id=record.id,
extract_metadata=True,
compute_hash=True,
)
assert enriched is False
session.expire_all()
assert session.get(AssetContent, content.id).hash is None
assert session.get(Asset, record.id).system_metadata is None, (
"size matching alone is not proof the file is unchanged; an mtime "
"disagreement alone must also discard the result"
)
def test_enrichment_lands_metadata_and_hash_from_one_stable_observation(
session, temp_dir: Path
):
path = temp_dir / "stable-both.bin"
path.write_bytes(b"stable content for both")
content, record = _create_unhashed_record(session, path)
with patch(
"app.assets.scanner.snapshot_hash", return_value=("both-digest", path.stat())
):
enriched = enrich_asset(
session,
file_path=str(path),
content_id=content.id,
record_id=record.id,
extract_metadata=True,
compute_hash=True,
)
assert enriched is True
session.expire_all()
assert session.get(AssetContent, content.id).hash == to_stored_hash("both-digest")
assert session.get(Asset, record.id).system_metadata == {
"filename": "stable-both.bin",
"file_path": str(path),
"format": "bin",
"content_type": "application/octet-stream",
"content_length": path.stat().st_size,
}
def test_enrichment_reports_no_progress_when_a_requested_hash_fails(
session, temp_dir: Path
):
path = temp_dir / "unreadable-for-hashing.bin"
path.write_bytes(b"metadata reads fine, hashing does not")
content, record = _create_unhashed_record(session, path)
with patch("app.assets.scanner.snapshot_hash", side_effect=PermissionError("denied")):
enriched = enrich_asset(
session,
file_path=str(path),
content_id=content.id,
record_id=record.id,
extract_metadata=True,
compute_hash=True,
)
assert enriched is False, (
"a record whose requested hash could not be computed has made no progress; "
"counting the metadata as progress keeps the record out of failed_ids, so the "
"enrich pass re-selects it forever"
)
session.expire_all()
assert session.get(AssetContent, content.id).hash is None
def test_off_mode_enrichment_still_lands_metadata_without_a_hash(
session, temp_dir: Path
):
path = temp_dir / "off-mode.bin"
path.write_bytes(b"metadata only")
content, record = _create_unhashed_record(session, path)
with patch("app.assets.scanner.snapshot_hash") as never_hashed:
enriched = enrich_asset(
session,
file_path=str(path),
content_id=content.id,
record_id=record.id,
extract_metadata=True,
compute_hash=False,
)
never_hashed.assert_not_called()
assert enriched is True
session.expire_all()
assert session.get(AssetContent, content.id).hash is None
assert session.get(Asset, record.id).system_metadata == {
"filename": "off-mode.bin",
"file_path": str(path),
"format": "bin",
"content_type": "application/octet-stream",
"content_length": path.stat().st_size,
}

View File

@@ -0,0 +1,118 @@
from unittest.mock import patch
from sqlalchemy import select
from app.assets.database.models import Asset, AssetContent
from app.assets.database.queries.records import create_content, mark_content_missing
from app.assets.helpers import to_stored_hash
from app.assets.services.ingest import create_from_hash
from app.assets.services.lookup import (
claim_qualified_content as _real_claim_qualified_content,
refresh_qualified_content as _real_refresh_qualified_content,
)
def test_create_from_hash_with_prefixed_hash_finds_existing_content(
mock_create_session, monkeypatch, temp_dir
):
digest = "a" * 64
path = temp_dir / "existing.bin"
path.write_bytes(b"existing bytes")
monkeypatch.setattr("app.assets.mode.hashing_enabled", lambda: True)
with mock_create_session() as session:
content = create_content(session, str(path), to_stored_hash(digest), path.stat().st_size)
content_id = content.id
session.commit()
result = create_from_hash(f"blake3:{digest}", "derived.bin")
assert result is not None
with mock_create_session() as session:
records = list(session.scalars(select(Asset)))
created_record = session.get(Asset, result.ref.id)
assert created_record is not None
assert created_record.content_id == content_id
assert [record.content_id for record in records] == [content_id]
def test_create_from_hash_with_bare_hash_also_works(
mock_create_session, monkeypatch, temp_dir
):
digest = "b" * 64
path = temp_dir / "existing.bin"
path.write_bytes(b"existing bytes")
monkeypatch.setattr("app.assets.mode.hashing_enabled", lambda: True)
with mock_create_session() as session:
content = create_content(session, str(path), to_stored_hash(digest), path.stat().st_size)
content_id = content.id
session.commit()
result = create_from_hash(digest, "derived.bin")
assert result is not None
with mock_create_session() as session:
records = list(session.scalars(select(Asset)))
created_record = session.get(Asset, result.ref.id)
assert created_record is not None
assert created_record.content_id == content_id
assert [record.content_id for record in records] == [content_id]
def _seed_live_content(mock_create_session, path, digest):
with mock_create_session() as session:
content = create_content(session, str(path), to_stored_hash(digest), path.stat().st_size)
content_id = content.id
session.commit()
return content_id
def test_content_retired_between_lookup_and_claim_mints_nothing(
mock_create_session, monkeypatch, temp_dir
):
digest = "c" * 64
path = temp_dir / "retired.bin"
path.write_bytes(b"retired bytes")
monkeypatch.setattr("app.assets.mode.hashing_enabled", lambda: True)
content_id = _seed_live_content(mock_create_session, path, digest)
def retire_then_claim(session, claimed_id, hash):
mark_content_missing(session, claimed_id)
session.commit()
return _real_claim_qualified_content(session, claimed_id, hash)
with patch("app.assets.services.ingest.claim_qualified_content", retire_then_claim):
result = create_from_hash(f"blake3:{digest}", "derived.bin")
assert result is None
with mock_create_session() as session:
assert list(session.scalars(select(Asset))) == []
retired = session.get(AssetContent, content_id)
assert retired is not None
assert retired.hash == to_stored_hash(digest)
assert retired.path == str(path)
def test_file_vanishing_between_claim_and_refresh_mints_nothing(
mock_create_session, monkeypatch, temp_dir
):
digest = "d" * 64
path = temp_dir / "vanishing.bin"
path.write_bytes(b"vanishing bytes")
monkeypatch.setattr("app.assets.mode.hashing_enabled", lambda: True)
content_id = _seed_live_content(mock_create_session, path, digest)
def delete_file_then_refresh(session, refreshed_id):
path.unlink()
return _real_refresh_qualified_content(session, refreshed_id)
with patch(
"app.assets.services.ingest.refresh_qualified_content", delete_file_then_refresh
):
result = create_from_hash(f"blake3:{digest}", "derived.bin")
assert result is None
with mock_create_session() as session:
assert list(session.scalars(select(Asset))) == []
assert session.get(AssetContent, content_id) is not None

View File

@@ -0,0 +1,58 @@
from unittest.mock import patch
from app.assets.database.models import AssetContent
from app.assets.services.hash_mode_state import (
clear_transition_queue,
drain_transition_queue,
enqueue_transition_work,
pending_transition_count,
read_stored_mode,
record_transition_intent,
write_stored_mode,
)
def test_absent_row_off_mode_no_transition(session):
with patch("app.assets.services.hash_mode_state._mode.hashing_enabled", return_value=False):
assert record_transition_intent(session) is None
assert read_stored_mode(session) == "off"
def test_empty_drain_keeps_off_mode_ready_for_a_later_on_transition(session):
write_stored_mode(session, "off")
drain_transition_queue(session)
with patch("app.assets.services.hash_mode_state._mode.hashing_enabled", return_value=True):
assert record_transition_intent(session) == "off_to_on"
def test_off_to_on_enqueues_null_rows(session):
session.add(AssetContent(path="/tmp/null", hash=None))
session.add(AssetContent(path="/tmp/hashed", hash="abc"))
write_stored_mode(session, "off")
with patch("app.assets.services.hash_mode_state._mode.hashing_enabled", return_value=True):
transition = record_transition_intent(session)
enqueue_transition_work(session, transition)
assert transition == "off_to_on"
assert pending_transition_count() == 2
assert read_stored_mode(session) == "off"
clear_transition_queue()
def test_on_to_off_freezes(session):
write_stored_mode(session, "on")
with patch("app.assets.services.hash_mode_state._mode.hashing_enabled", return_value=False):
transition = record_transition_intent(session)
assert transition == "on_to_off"
assert read_stored_mode(session) == "off"
assert pending_transition_count() == 0
def test_mode_stays_off_during_transition(session):
session.add(AssetContent(path="/tmp/null", hash=None))
write_stored_mode(session, "off")
with patch("app.assets.services.hash_mode_state._mode.hashing_enabled", return_value=True):
transition = record_transition_intent(session)
enqueue_transition_work(session, transition)
assert read_stored_mode(session) == "off"

View File

@@ -1,568 +1,23 @@
"""Tests for ingest services."""
from contextlib import contextmanager
from pathlib import Path
from unittest.mock import patch
import pytest
from PIL import Image
from sqlalchemy.orm import Session as SASession, Session
from sqlalchemy import select
from app.assets.database.models import Asset, AssetReference, AssetReferenceTag, Tag
from app.assets.database.queries import get_reference_tags
from app.assets.helpers import get_utc_now
from app.assets.services.ingest import (
_ingest_file_from_path,
_register_existing_asset,
ingest_existing_file,
)
import folder_paths
from app.assets.database.models import Asset, AssetContent
from app.assets.services.ingest import register_executed_output
def _make_png(path: Path, size: tuple[int, int]) -> Path:
Image.new("RGB", size, color=(80, 120, 200)).save(path, format="PNG")
return path
def test_output_registration_creates_separate_record_and_content_rows(mock_create_session) -> None:
path = Path(folder_paths.get_output_directory()) / "result-b-schema.txt"
path.parent.mkdir(parents=True, exist_ok=True)
path.write_text("result")
registered = register_executed_output(str(path), job_id="job")
class TestIngestFileFromPath:
def test_creates_asset_and_reference(self, mock_create_session, temp_dir: Path, session: Session):
file_path = temp_dir / "test_file.bin"
file_path.write_bytes(b"test content")
result = _ingest_file_from_path(
abs_path=str(file_path),
asset_hash="blake3:abc123",
size_bytes=12,
mtime_ns=1234567890000000000,
mime_type="application/octet-stream",
)
assert result.asset_created is True
assert result.ref_created is True
assert result.reference_id is not None
# Verify DB state
assets = session.query(Asset).all()
assert len(assets) == 1
assert assets[0].hash == "blake3:abc123"
refs = session.query(AssetReference).all()
assert len(refs) == 1
assert refs[0].file_path == str(file_path)
def test_creates_reference_when_name_provided(self, mock_create_session, temp_dir: Path, session: Session):
file_path = temp_dir / "model.safetensors"
file_path.write_bytes(b"model data")
result = _ingest_file_from_path(
abs_path=str(file_path),
asset_hash="blake3:def456",
size_bytes=10,
mtime_ns=1234567890000000000,
mime_type="application/octet-stream",
info_name="My Model",
owner_id="user1",
)
assert result.asset_created is True
assert result.reference_id is not None
ref = session.query(AssetReference).first()
assert ref is not None
assert ref.name == "My Model"
assert ref.owner_id == "user1"
def test_creates_tags_when_provided(self, mock_create_session, temp_dir: Path, session: Session):
file_path = temp_dir / "tagged.bin"
file_path.write_bytes(b"data")
result = _ingest_file_from_path(
abs_path=str(file_path),
asset_hash="blake3:ghi789",
size_bytes=4,
mtime_ns=1234567890000000000,
info_name="Tagged Asset",
tags=["models", "checkpoints"],
)
assert result.reference_id is not None
# Verify tags were created and linked
tags = session.query(Tag).all()
tag_names = {t.name for t in tags}
assert "models" in tag_names
assert "checkpoints" in tag_names
ref_tags = get_reference_tags(session, reference_id=result.reference_id)
assert set(ref_tags) == {"models", "checkpoints"}
def test_path_derived_tags_use_automatic_origin(
self, mock_create_session, temp_dir: Path, session: Session
):
input_dir = temp_dir / "input"
output_dir = temp_dir / "output"
temp_root = temp_dir / "temp"
for directory in (input_dir, output_dir, temp_root):
directory.mkdir()
file_path = input_dir / "pasted" / "tagged.png"
file_path.parent.mkdir()
file_path.write_bytes(b"data")
with (
patch("app.assets.services.path_utils.folder_paths") as mock_fp,
patch(
"app.assets.services.path_utils.get_comfy_models_folders",
return_value=[],
),
):
mock_fp.get_input_directory.return_value = str(input_dir)
mock_fp.get_output_directory.return_value = str(output_dir)
mock_fp.get_temp_directory.return_value = str(temp_root)
result = _ingest_file_from_path(
abs_path=str(file_path),
asset_hash="blake3:pathorigin",
size_bytes=4,
mtime_ns=1234567890000000000,
info_name="Tagged Asset",
tags=["input", "manual-label"],
)
assert result.reference_id is not None
links = session.query(AssetReferenceTag).filter_by(
asset_reference_id=result.reference_id
)
origin_by_tag = {link.tag_name: link.origin for link in links}
assert origin_by_tag["input"] == "automatic"
assert origin_by_tag["pasted"] == "automatic"
assert origin_by_tag["manual-label"] == "manual"
def test_idempotent_upsert(self, mock_create_session, temp_dir: Path, session: Session):
file_path = temp_dir / "dup.bin"
file_path.write_bytes(b"content")
# First ingest
r1 = _ingest_file_from_path(
abs_path=str(file_path),
asset_hash="blake3:repeat",
size_bytes=7,
mtime_ns=1234567890000000000,
)
assert r1.asset_created is True
# Second ingest with same hash - should update, not create
r2 = _ingest_file_from_path(
abs_path=str(file_path),
asset_hash="blake3:repeat",
size_bytes=7,
mtime_ns=1234567890000000001, # different mtime
)
assert r2.asset_created is False
assert r2.ref_created is False
assert r2.ref_updated is True
# Still only one asset
assets = session.query(Asset).all()
assert len(assets) == 1
def test_validates_preview_id(self, mock_create_session, temp_dir: Path, session: Session):
file_path = temp_dir / "with_preview.bin"
file_path.write_bytes(b"data")
# Create a preview asset and reference
preview_asset = Asset(hash="blake3:preview", size_bytes=100)
session.add(preview_asset)
session.flush()
from app.assets.helpers import get_utc_now
now = get_utc_now()
preview_ref = AssetReference(
asset_id=preview_asset.id, name="preview.png", owner_id="",
created_at=now, updated_at=now, last_access_time=now,
)
session.add(preview_ref)
session.commit()
preview_id = preview_ref.id
result = _ingest_file_from_path(
abs_path=str(file_path),
asset_hash="blake3:main",
size_bytes=4,
mtime_ns=1234567890000000000,
info_name="With Preview",
preview_id=preview_id,
)
assert result.reference_id is not None
ref = session.query(AssetReference).filter_by(id=result.reference_id).first()
assert ref.preview_id == preview_id
def test_invalid_preview_id_is_cleared(self, mock_create_session, temp_dir: Path, session: Session):
file_path = temp_dir / "bad_preview.bin"
file_path.write_bytes(b"data")
result = _ingest_file_from_path(
abs_path=str(file_path),
asset_hash="blake3:badpreview",
size_bytes=4,
mtime_ns=1234567890000000000,
info_name="Bad Preview",
preview_id="nonexistent-uuid",
)
assert result.reference_id is not None
ref = session.query(AssetReference).filter_by(id=result.reference_id).first()
assert ref.preview_id is None
class TestRegisterExistingAsset:
def test_creates_reference_for_existing_asset(self, mock_create_session, session: Session):
# Create existing asset
asset = Asset(hash="blake3:existing", size_bytes=1024, mime_type="image/png")
session.add(asset)
session.commit()
result = _register_existing_asset(
asset_hash="blake3:existing",
name="Registered Asset",
user_metadata={"key": "value"},
tags=["models"],
)
assert result.created is True
assert "models" in result.tags
# Verify by re-fetching from DB
session.expire_all()
refs = session.query(AssetReference).filter_by(name="Registered Asset").all()
assert len(refs) == 1
def test_creates_new_reference_even_with_same_name(self, mock_create_session, session: Session):
# Create asset and reference
asset = Asset(hash="blake3:withref", size_bytes=512)
session.add(asset)
session.flush()
from app.assets.helpers import get_utc_now
ref = AssetReference(
owner_id="",
name="Existing Ref",
asset_id=asset.id,
created_at=get_utc_now(),
updated_at=get_utc_now(),
last_access_time=get_utc_now(),
)
session.add(ref)
session.flush()
ref_id = ref.id
session.commit()
result = _register_existing_asset(
asset_hash="blake3:withref",
name="Existing Ref",
owner_id="",
)
# Multiple files with same name are allowed
assert result.created is True
# Verify two AssetReferences exist for this name
session.expire_all()
refs = session.query(AssetReference).filter_by(name="Existing Ref").all()
assert len(refs) == 2
assert ref_id in [r.id for r in refs]
def test_raises_for_nonexistent_hash(self, mock_create_session):
with pytest.raises(ValueError, match="No asset with hash"):
_register_existing_asset(
asset_hash="blake3:doesnotexist",
name="Fail",
)
def test_applies_tags_to_new_reference(self, mock_create_session, session: Session):
asset = Asset(hash="blake3:tagged", size_bytes=256)
session.add(asset)
session.commit()
result = _register_existing_asset(
asset_hash="blake3:tagged",
name="Tagged Ref",
tags=["alpha", "beta"],
)
assert result.created is True
assert set(result.tags) == {"alpha", "beta"}
class TestIngestExistingFileTagFK:
"""Regression: ingest_existing_file must seed Tag rows before inserting
AssetReferenceTag rows, otherwise FK enforcement raises IntegrityError."""
def test_creates_tag_rows_before_reference_tags(self, db_engine_fk, temp_dir: Path):
"""With PRAGMA foreign_keys=ON, tags must exist in the tags table
before they can be referenced in asset_reference_tags."""
@contextmanager
def _create_session():
with SASession(db_engine_fk) as sess:
yield sess
file_path = temp_dir / "output.png"
file_path.write_bytes(b"image data")
with patch("app.assets.services.ingest.create_session", _create_session), \
patch(
"app.assets.services.ingest.get_name_and_tags_from_asset_path",
return_value=("output.png", ["output"]),
):
result = ingest_existing_file(
abs_path=str(file_path),
extra_tags=["my-job"],
)
assert result is True
with SASession(db_engine_fk) as sess:
tag_names = {t.name for t in sess.query(Tag).all()}
assert "output" in tag_names
assert "my-job" in tag_names
ref_tags = sess.query(AssetReferenceTag).all()
ref_tag_names = {rt.tag_name for rt in ref_tags}
assert "output" in ref_tag_names
class TestIngestExistingFileLoaderPath:
"""Outputs saved into a subfolder must persist the subfolder-qualified
loader path, not the bare basename (regression: spec["fname"] was
os.path.basename)."""
def test_subfoldered_output_persists_relative_loader_path(
self, mock_create_session, temp_dir: Path, session: Session
):
input_dir = temp_dir / "input"
output_dir = temp_dir / "output"
temp_root = temp_dir / "temp"
for directory in (input_dir, output_dir, temp_root):
directory.mkdir()
file_path = output_dir / "sub" / "img_00001_.png"
file_path.parent.mkdir()
file_path.write_bytes(b"image data")
with (
patch("app.assets.services.path_utils.folder_paths") as mock_fp,
patch(
"app.assets.services.path_utils.get_comfy_models_folders",
return_value=[],
),
):
mock_fp.get_input_directory.return_value = str(input_dir)
mock_fp.get_output_directory.return_value = str(output_dir)
mock_fp.get_temp_directory.return_value = str(temp_root)
assert ingest_existing_file(abs_path=str(file_path)) is True
ref = (
session.query(AssetReference)
.filter_by(file_path=str(file_path))
.one()
)
assert ref.loader_path == "sub/img_00001_.png"
assert (ref.user_metadata or {}).get("filename") == "sub/img_00001_.png"
class TestIngestImageDimensions:
"""system_metadata should carry {kind, width, height} for image assets."""
def test_image_asset_emits_dimensions(
self, mock_create_session, temp_dir: Path, session: Session
):
f = _make_png(temp_dir / "shot.png", (640, 480))
result = _ingest_file_from_path(
abs_path=str(f),
asset_hash="blake3:img1",
size_bytes=f.stat().st_size,
mtime_ns=1234567890000000000,
mime_type="image/png",
)
ref = session.query(AssetReference).filter_by(id=result.reference_id).first()
assert ref.system_metadata == {
"kind": "image",
"width": 640,
"height": 480,
}
def test_non_image_asset_leaves_system_metadata_empty(
self, mock_create_session, temp_dir: Path, session: Session
):
f = temp_dir / "model.safetensors"
f.write_bytes(b"not an image")
result = _ingest_file_from_path(
abs_path=str(f),
asset_hash="blake3:safetensors1",
size_bytes=f.stat().st_size,
mtime_ns=1234567890000000000,
mime_type="application/octet-stream",
)
ref = session.query(AssetReference).filter_by(id=result.reference_id).first()
assert ref.system_metadata in (None, {})
def test_preserves_existing_system_metadata_keys(
self, mock_create_session, temp_dir: Path, session: Session
):
f = _make_png(temp_dir / "annotated.png", (100, 200))
# First pass populates a sentinel system_metadata key (simulating prior
# enricher write).
result = _ingest_file_from_path(
abs_path=str(f),
asset_hash="blake3:img-merge",
size_bytes=f.stat().st_size,
mtime_ns=1234567890000000000,
mime_type="image/png",
)
ref = session.query(AssetReference).filter_by(id=result.reference_id).first()
ref.system_metadata = {**(ref.system_metadata or {}), "source_url": "https://example/x.png"}
session.commit()
# Second pass with the same path triggers the merge code path again.
_ingest_file_from_path(
abs_path=str(f),
asset_hash="blake3:img-merge",
size_bytes=f.stat().st_size,
mtime_ns=1234567890000000001,
mime_type="image/png",
)
session.refresh(ref)
assert ref.system_metadata["kind"] == "image"
assert ref.system_metadata["width"] == 100
assert ref.system_metadata["height"] == 200
assert ref.system_metadata["source_url"] == "https://example/x.png"
class TestRegisterExistingAssetBackfill:
"""The from-hash path back-fills dimensions from a sibling reference."""
def _add_reference(
self,
session: Session,
asset: Asset,
name: str,
system_metadata: dict | None = None,
) -> AssetReference:
now = get_utc_now()
ref = AssetReference(
asset_id=asset.id,
name=name,
owner_id="",
created_at=now,
updated_at=now,
last_access_time=now,
system_metadata=system_metadata or {},
)
session.add(ref)
session.flush()
return ref
def test_backfills_dimensions_from_sibling_image_reference(
self, mock_create_session, session: Session
):
asset = Asset(hash="blake3:shared", size_bytes=2048, mime_type="image/png")
session.add(asset)
session.flush()
self._add_reference(
session,
asset,
name="original.png",
system_metadata={"kind": "image", "width": 800, "height": 600},
)
session.commit()
result = _register_existing_asset(
asset_hash="blake3:shared",
name="from_hash.png",
owner_id="user-x",
)
ref = session.query(AssetReference).filter_by(id=result.ref.id).first()
assert ref.system_metadata.get("kind") == "image"
assert ref.system_metadata.get("width") == 800
assert ref.system_metadata.get("height") == 600
def test_no_backfill_when_sibling_has_no_image_metadata(
self, mock_create_session, session: Session
):
asset = Asset(hash="blake3:nodims", size_bytes=2048, mime_type="image/png")
session.add(asset)
session.flush()
self._add_reference(
session,
asset,
name="original.png",
system_metadata={"base_model": "flux"}, # no kind=image
)
session.commit()
result = _register_existing_asset(
asset_hash="blake3:nodims",
name="from_hash.png",
owner_id="user-x",
)
ref = session.query(AssetReference).filter_by(id=result.ref.id).first()
meta = ref.system_metadata or {}
assert "kind" not in meta
assert "width" not in meta
assert "height" not in meta
def test_no_backfill_when_no_sibling_exists(
self, mock_create_session, session: Session
):
asset = Asset(hash="blake3:lonely", size_bytes=1024, mime_type="image/png")
session.add(asset)
session.commit()
result = _register_existing_asset(
asset_hash="blake3:lonely",
name="solo.png",
owner_id="user-x",
)
ref = session.query(AssetReference).filter_by(id=result.ref.id).first()
assert ref.system_metadata in (None, {})
def test_backfill_preserves_caller_supplied_keys(
self, mock_create_session, session: Session
):
asset = Asset(hash="blake3:preserve", size_bytes=2048, mime_type="image/png")
session.add(asset)
session.flush()
self._add_reference(
session,
asset,
name="original.png",
system_metadata={"kind": "image", "width": 1024, "height": 768},
)
session.commit()
# Simulate a from-hash path where the new reference already carries
# some system_metadata (e.g. a download-provenance source_url written
# by an earlier step). The back-fill must merge dim keys without
# clobbering existing keys.
result = _register_existing_asset(
asset_hash="blake3:preserve",
name="from_hash.png",
owner_id="user-x",
)
ref = session.query(AssetReference).filter_by(id=result.ref.id).first()
# Seed a sentinel key and re-run back-fill via a second register call
# to exercise the merge path with pre-existing data.
ref.system_metadata = {**(ref.system_metadata or {}), "source_url": "https://example/p"}
session.commit()
assert ref.system_metadata.get("source_url") == "https://example/p"
assert ref.system_metadata.get("kind") == "image"
assert ref.system_metadata.get("width") == 1024
assert ref.system_metadata.get("height") == 768
with mock_create_session() as session:
record = session.get(Asset, registered.id)
content = session.scalar(select(AssetContent).where(AssetContent.id == registered.content_id))
assert record is not None
assert content is not None
assert record.content_id == content.id
assert record.job_id == "job"

View File

@@ -0,0 +1,82 @@
import os
import pytest
from sqlalchemy import select
import app.assets.mode as mode_module
import folder_paths
from app.assets.database.models import Asset, AssetContent
from app.assets.services.ingest import register_executed_output
@pytest.fixture(autouse=True)
def hashing_off():
class FakeArgs:
enable_asset_hashing = False
mode_module.init(FakeArgs())
yield
mode_module.init(None)
def test_new_path_save_off_mode_hash_null(mock_create_session):
output_dir = folder_paths.get_output_directory()
os.makedirs(output_dir, exist_ok=True)
f = os.path.join(output_dir, "test_ingest_b_new.png")
with open(f, "wb") as fh:
fh.write(b"pixels")
try:
record = register_executed_output(f, job_id="job1")
record_id = record.id
with mock_create_session() as session:
content = session.execute(
select(AssetContent).where(AssetContent.path == os.path.abspath(f))
).scalar_one()
assert content.hash is None
assert content.is_missing is False
asset = session.execute(
select(Asset).where(Asset.id == record_id)
).scalar_one()
assert asset.job_id == "job1"
finally:
if os.path.exists(f):
os.unlink(f)
def test_overwrite_at_live_path_marks_old_missing(mock_create_session):
output_dir = folder_paths.get_output_directory()
os.makedirs(output_dir, exist_ok=True)
f = os.path.join(output_dir, "test_ingest_b_overwrite.png")
try:
with open(f, "wb") as fh:
fh.write(b"v1")
r1 = register_executed_output(f, job_id="job1")
old_content_id = r1.content_id
old_record_id = r1.id
with open(f, "wb") as fh:
fh.write(b"v2")
r2 = register_executed_output(f, job_id="job2")
new_record_id = r2.id
with mock_create_session() as session:
old_content = session.execute(
select(AssetContent).where(AssetContent.id == old_content_id)
).scalar_one()
assert old_content.is_missing is True, "Old content should be marked missing"
old_record = session.execute(
select(Asset).where(Asset.id == old_record_id)
).scalar_one()
assert old_record.job_id == "job1", "Old record's job_id must not be mutated"
new_record = session.execute(
select(Asset).where(Asset.id == new_record_id)
).scalar_one()
assert new_record.job_id == "job2"
assert new_record_id != old_record_id
finally:
if os.path.exists(f):
os.unlink(f)

View File

@@ -0,0 +1,168 @@
import os
import uuid
import pytest
from sqlalchemy import func, select
from sqlalchemy.orm import Session
import app.assets.mode as mode_module
import folder_paths
from app.assets.database.models import Asset, AssetContent
from app.assets.database.queries import create_record as create_record_query
from app.assets.scanner import SeedAssetSpec, seed_asset_specs
from app.assets.services import ingest
from app.assets.services.ingest import register_file_in_place, upload_from_temp_path
@pytest.fixture
def hashing_off():
class FakeArgs:
enable_asset_hashing = False
mode_module.init(FakeArgs())
yield
mode_module.init(None)
def _write_temp(content: bytes) -> str:
uploads_root = os.path.join(
folder_paths.get_temp_directory(), "uploads", uuid.uuid4().hex
)
os.makedirs(uploads_root, exist_ok=True)
path = os.path.join(uploads_root, ".upload.part")
with open(path, "wb") as file:
file.write(content)
return path
def _raise_create_record(*_args, **_kwargs):
raise RuntimeError("forced create_record failure")
def _live_content_count(session: Session) -> int:
return session.scalar(
select(func.count())
.select_from(AssetContent)
.where(AssetContent.is_missing.is_(False))
)
def test_upload_record_failure_discards_newly_inserted_content(
mock_create_session, hashing_off, monkeypatch: pytest.MonkeyPatch
) -> None:
content = b"orphan-upload-" + uuid.uuid4().bytes
temp_path = _write_temp(content)
monkeypatch.setattr(ingest, "create_record", _raise_create_record)
with pytest.raises(RuntimeError, match="forced create_record failure"):
upload_from_temp_path(
temp_path=temp_path,
name="orphan.bin",
tags=["output"],
client_filename="orphan.bin",
)
with mock_create_session() as session:
assert _live_content_count(session) == 0
def test_register_file_in_place_discards_content_on_record_failure(
mock_create_session, hashing_off, monkeypatch: pytest.MonkeyPatch
) -> None:
output_dir = folder_paths.get_output_directory()
os.makedirs(output_dir, exist_ok=True)
path = os.path.join(output_dir, f"orphan_inplace_{uuid.uuid4().hex}.png")
with open(path, "wb") as file:
file.write(b"orphan-inplace-" + uuid.uuid4().bytes)
monkeypatch.setattr(ingest, "create_record", _raise_create_record)
try:
with pytest.raises(RuntimeError, match="forced create_record failure"):
register_file_in_place(abs_path=path, name="orphan_inplace.png", tags=["output"])
with mock_create_session() as session:
assert _live_content_count(session) == 0
finally:
if os.path.exists(path):
os.unlink(path)
def _seed_spec(path: str, size_bytes: int, mtime_ns: int, name: str) -> SeedAssetSpec:
return {
"abs_path": path,
"size_bytes": size_bytes,
"mtime_ns": mtime_ns,
"info_name": name,
"tags": ["input"],
"fname": name,
"metadata": None,
"mime_type": None,
"job_id": None,
}
def _content_at(session: Session, path: str) -> AssetContent | None:
return session.scalar(select(AssetContent).where(AssetContent.path == path))
def _reference_count(session: Session, content_id: str) -> int:
return session.scalar(
select(func.count()).select_from(Asset).where(Asset.content_id == content_id)
)
def _orphaned_content_paths(session: Session) -> list[str]:
return [
content.path
for content in session.scalars(select(AssetContent))
if _reference_count(session, content.id) == 0
]
def test_seed_asset_specs_orphans_nothing_and_keeps_earlier_specs_on_record_failure(
session: Session, tmp_path, monkeypatch: pytest.MonkeyPatch
) -> None:
specs: list[SeedAssetSpec] = []
paths: dict[str, str] = {}
fail_name = "vanished.bin"
for name in ("first.bin", fail_name, "last.bin"):
file_path = tmp_path / name
file_path.write_bytes(name.encode())
stat_result = file_path.stat()
paths[name] = str(file_path)
specs.append(
_seed_spec(str(file_path), stat_result.st_size, stat_result.st_mtime_ns, name)
)
attempted: list[str] = []
def _create_record_or_raise(session_arg, content_id, name, *args, **kwargs):
attempted.append(name)
if name == fail_name:
raise RuntimeError("forced create_record failure")
return create_record_query(session_arg, content_id, name, *args, **kwargs)
monkeypatch.setattr("app.assets.scanner.create_record", _create_record_or_raise)
with pytest.raises(RuntimeError, match="forced create_record failure"):
seed_asset_specs(session, specs)
session.rollback()
assert _content_at(session, paths[fail_name]) is None, (
"the failed spec's content must not outlive the record that would have referenced it"
)
assert session.scalar(select(Asset).where(Asset.name == fail_name)) is None
assert _orphaned_content_paths(session) == []
survivor = _content_at(session, paths["first.bin"])
assert survivor is not None, (
"a spec that already succeeded must not be retroactively erased by a later "
"unrelated failure; its savepoint was released before that failure happened"
)
assert session.scalar(select(Asset).where(Asset.name == "first.bin")) is not None
assert _reference_count(session, survivor.id) == 1
assert attempted == ["first.bin", fail_name]
assert _content_at(session, paths["last.bin"]) is None, (
"the raise aborts the loop, so the spec after the failed one is never attempted"
)

View File

@@ -0,0 +1,333 @@
import logging
import os
import tempfile
from contextlib import contextmanager
from pathlib import Path
from unittest.mock import patch
import pytest
from sqlalchemy import create_engine, select
from sqlalchemy.orm import Session
from sqlalchemy.orm import Session as SASession
from app.assets import lifecycle
from app.assets.database.models import Asset, AssetContent, Base
from app.assets.database.queries.records import create_content, create_record
from app.assets.lifecycle import (
cleanup_temp_filesystem,
get_excluded_scan_roots,
run_asset_shutdown_cleanup,
run_asset_startup,
run_startup,
wipe_temp_db_rows,
)
from app.assets.scanner import get_temp_prefixes, sync_temp_references_safely
from app.assets.scanner_changes import is_path_under_prefixes
from app.assets.seeder import asset_seeder
from app.assets.services import hash_mode_state
from app.assets.services.hash_mode_state import clear_transition_queue, write_stored_mode
from .path_prefix_cases import expected_prefix_case_paths, prefix_case_paths
@pytest.fixture(autouse=True)
def autoclean_unit_test_assets():
lifecycle._excluded_scan_roots.clear()
yield
lifecycle._excluded_scan_roots.clear()
@pytest.fixture(autouse=True)
def autoclean_transition_queue():
clear_transition_queue()
lifecycle._hash_mode_transition = None
yield
clear_transition_queue()
lifecycle._hash_mode_transition = None
@pytest.fixture
def session():
engine = create_engine("sqlite:///:memory:", connect_args={"check_same_thread": False})
Base.metadata.create_all(engine)
with Session(engine) as sess:
yield sess
@pytest.fixture
def mock_create_session(session):
engine = session.bind
@contextmanager
def _create_session():
with SASession(engine) as sess:
yield sess
with patch("app.assets.lifecycle.create_session", _create_session), \
patch("app.database.db.create_session", _create_session):
yield _create_session
@pytest.fixture
def comfy_dirs():
with tempfile.TemporaryDirectory() as base:
temp_dir = Path(base) / "temp"
temp_dir.mkdir()
with patch("folder_paths.get_temp_directory", return_value=str(temp_dir)):
yield temp_dir
def _seed_temp_rows(session: Session, temp_dir: Path) -> tuple[str, str]:
temp_file = temp_dir / "preview.png"
temp_file.write_bytes(b"\x00" * 10)
content = create_content(session, path=str(temp_file))
record = create_record(session, content_id=content.id, name="preview.png")
session.commit()
return record.id, content.id
def test_startup_order_wipe_before_rmtree_before_seeder(mock_create_session):
calls: list[str] = []
def _wipe(session):
calls.append("wipe")
return 0, 0
with (
patch("app.assets.lifecycle.wipe_temp_db_rows", side_effect=_wipe),
patch("app.assets.lifecycle.cleanup_temp_filesystem", side_effect=lambda: calls.append("rmtree") or True),
patch("app.assets.lifecycle.enqueue_mode_transition_work", side_effect=lambda: calls.append("enqueue")),
patch("app.assets.lifecycle.start_asset_seeder", side_effect=lambda: calls.append("seeder") or True),
):
run_asset_startup()
assert calls == ["wipe", "rmtree", "enqueue", "seeder"]
@pytest.mark.parametrize("wipe_fails", [False, True], ids=["wipe_succeeds", "wipe_fails"])
def test_startup_defers_transition_drain_to_the_seeder(
session, mock_create_session, tmp_path, monkeypatch, wipe_fails
):
asset_file = tmp_path / "model.safetensors"
asset_file.write_bytes(b"weights")
stat = asset_file.stat()
create_content(session, path=str(asset_file), size_bytes=stat.st_size, mtime_ns=stat.st_mtime_ns)
write_stored_mode(session, "off")
session.commit()
monkeypatch.setattr(hash_mode_state._mode, "hashing_enabled", lambda: True)
wipe_side_effect = RuntimeError("db wipe failed") if wipe_fails else (lambda _session: (0, 0))
with (
patch("app.assets.lifecycle.wipe_temp_db_rows", side_effect=wipe_side_effect),
patch("app.assets.lifecycle.cleanup_temp_filesystem", return_value=True),
patch("app.assets.lifecycle.start_asset_seeder", return_value=True),
patch.object(hash_mode_state, "drain_transition_queue") as drain_spy,
):
lifecycle.record_hash_mode_transition_intent()
run_asset_startup()
drain_spy.assert_not_called()
assert hash_mode_state.pending_transition_count() > 0
assert hash_mode_state.read_stored_mode(session) == "off"
def test_db_wipe_failure_skips_rmtree_and_continues(mock_create_session):
with (
patch("app.assets.lifecycle.wipe_temp_db_rows", side_effect=RuntimeError("db wipe failed")),
patch("app.assets.lifecycle.cleanup_temp_filesystem") as cleanup_mock,
patch("app.assets.lifecycle.start_asset_seeder", return_value=True) as seeder_mock,
):
run_asset_startup()
cleanup_mock.assert_not_called()
seeder_mock.assert_called_once()
def test_run_startup_disabled_sweeps_temp_filesystem_without_db_work(comfy_dirs):
stale = comfy_dirs / "stale.png"
stale.write_bytes(b"\x00" * 10)
assert stale.exists()
with (
patch("app.assets.lifecycle.run_asset_startup") as asset_startup_mock,
patch("app.assets.lifecycle.wipe_temp_db_rows") as wipe_mock,
patch("app.assets.lifecycle.start_asset_seeder") as seeder_mock,
):
run_startup(enable_assets=False)
assert not stale.exists()
asset_startup_mock.assert_not_called()
wipe_mock.assert_not_called()
seeder_mock.assert_not_called()
def test_run_startup_enabled_delegates_to_asset_startup_not_bare_sweep():
with (
patch("app.assets.lifecycle.run_asset_startup") as asset_startup_mock,
patch("app.assets.lifecycle.cleanup_temp_filesystem") as cleanup_mock,
):
run_startup(enable_assets=True)
asset_startup_mock.assert_called_once_with()
cleanup_mock.assert_not_called()
def test_run_startup_logs_and_absorbs_disabled_filesystem_failure(caplog):
with patch("app.assets.lifecycle.cleanup_temp_filesystem", side_effect=RuntimeError("filesystem failure")):
run_startup(enable_assets=False)
assert "Asset startup maintenance failed" in caplog.text
def test_rmtree_failure_excludes_temp_from_scan(session, comfy_dirs, mock_create_session):
record_id, content_id = _seed_temp_rows(session, comfy_dirs)
wipe_temp_db_rows(session)
session.commit()
assert session.get(Asset, record_id) is None
with patch("app.assets.lifecycle.shutil.rmtree", side_effect=OSError("busy")):
assert cleanup_temp_filesystem() is False
assert str(comfy_dirs) in get_excluded_scan_roots()
assert get_temp_prefixes() == []
residual = comfy_dirs / "leftover.png"
residual.write_bytes(b"\x00" * 10)
with patch("app.assets.scanner.create_session", mock_create_session):
sync_temp_references_safely()
assert session.scalars(select(Asset)).all() == []
def test_shutdown_skips_cleanup_when_seeder_join_times_out(session, comfy_dirs, mock_create_session, caplog):
record_id, _ = _seed_temp_rows(session, comfy_dirs)
with patch("app.assets.lifecycle.wipe_temp_db_rows") as wipe_mock:
with patch.object(asset_seeder, "shutdown", return_value=False):
joined = asset_seeder.shutdown()
if joined:
run_asset_shutdown_cleanup()
wipe_mock.assert_not_called()
assert session.get(Asset, record_id) is not None
def test_seeder_shutdown_logs_only_when_thread_does_not_exit(caplog: pytest.LogCaptureFixture) -> None:
seeder = type(asset_seeder)()
with patch.object(seeder, "cancel"), patch.object(seeder, "wait", return_value=False), caplog.at_level(logging.WARNING):
assert seeder.shutdown(timeout=2.5) is False
assert caplog.messages == ["Asset seeder thread did not exit within 2.5s"]
assert "skipping temp cleanup" not in caplog.text
def test_shutdown_cleanup_wipes_rows_then_rmtree(session, comfy_dirs, mock_create_session):
record_id, content_id = _seed_temp_rows(session, comfy_dirs)
calls: list[str] = []
real_wipe = wipe_temp_db_rows
def tracking_wipe(sess):
calls.append("wipe")
return real_wipe(sess)
with (
patch("app.assets.lifecycle.wipe_temp_db_rows", side_effect=tracking_wipe),
patch("app.assets.lifecycle.cleanup_temp_filesystem", side_effect=lambda: calls.append("rmtree") or True),
):
run_asset_shutdown_cleanup()
assert calls == ["wipe", "rmtree"]
assert session.get(Asset, record_id) is None
assert session.get(AssetContent, content_id) is None
def test_run_shutdown_without_session_sweeps_temp_filesystem(comfy_dirs):
stale = comfy_dirs / "stale.png"
stale.write_bytes(b"\x00" * 10)
with (
patch("app.assets.lifecycle.can_create_session", return_value=False),
patch("app.assets.lifecycle.wipe_temp_db_rows") as wipe_mock,
):
lifecycle.run_shutdown()
assert not stale.exists()
wipe_mock.assert_not_called()
def test_wipe_temp_db_rows_deletes_records_before_content(session, comfy_dirs):
record_id, content_id = _seed_temp_rows(session, comfy_dirs)
records_deleted, contents_deleted = wipe_temp_db_rows(session)
session.commit()
assert records_deleted == 1
assert contents_deleted == 1
assert session.get(Asset, record_id) is None
assert session.get(AssetContent, content_id) is None
def test_wipe_temp_db_rows_preserves_non_temp_rows(session, comfy_dirs):
temp_record_id, temp_content_id = _seed_temp_rows(session, comfy_dirs)
keep_path = str(comfy_dirs) + "-sibling" + os.sep + "keep.safetensors"
keep_content = create_content(session, path=keep_path)
keep_record = create_record(session, content_id=keep_content.id, name="keep.safetensors")
session.commit()
records_deleted, contents_deleted = wipe_temp_db_rows(session)
session.commit()
assert (records_deleted, contents_deleted) == (1, 1)
assert session.get(Asset, temp_record_id) is None
assert session.get(AssetContent, temp_content_id) is None
assert session.get(Asset, keep_record.id) is not None
assert session.get(AssetContent, keep_content.id) is not None
def _seed_paths(session: Session, paths: list[str]) -> list[str]:
stored: list[str] = []
for index, path in enumerate(paths):
content = create_content(session, path=path)
create_record(session, content_id=content.id, name=f"case-{index}.png")
stored.append(content.path)
session.commit()
return stored
def _surviving_paths(session: Session) -> set[str]:
return set(session.scalars(select(AssetContent.path)).all())
def test_wipe_deletes_exactly_what_is_path_under_prefixes_accepts(session, comfy_dirs):
temp_root = str(comfy_dirs)
stored = _seed_paths(session, [path for path, _ in prefix_case_paths(temp_root)])
records_deleted, contents_deleted = wipe_temp_db_rows(session)
session.commit()
wiped = expected_prefix_case_paths(temp_root)
assert _surviving_paths(session) == set(stored) - wiped
assert (records_deleted, contents_deleted) == (len(wiped), len(wiped))
assert wiped and wiped != set(stored)
case_different = os.path.join(
os.path.dirname(temp_root), os.path.basename(temp_root).upper(), "case.png"
)
assert case_different in _surviving_paths(session)
assert os.path.join(os.path.dirname(temp_root), "escaped.png") in _surviving_paths(session)
def test_wipe_with_metacharacter_temp_root_matches_only_literal_children(session):
with tempfile.TemporaryDirectory() as base:
temp_root = os.path.join(base, "a_b%c*d?e[f")
inside = os.path.join(temp_root, "preview.png")
decoy = os.path.join(base, "aXbYcZdWeQf", "keep.safetensors")
_seed_paths(session, [inside, decoy])
with patch("folder_paths.get_temp_directory", return_value=temp_root):
wipe_temp_db_rows(session)
session.commit()
assert is_path_under_prefixes(decoy, [temp_root]) is False
assert _surviving_paths(session) == {decoy}

View File

@@ -0,0 +1,295 @@
import dataclasses
import threading
from collections.abc import Callable, Generator, Iterator
from contextlib import AbstractContextManager, contextmanager
from pathlib import Path
from typing import Protocol
from unittest.mock import MagicMock, Mock, call
import folder_paths
import pytest
from sqlalchemy import create_engine, select
from sqlalchemy.orm import Session, Session as SASession
from sqlalchemy.pool import StaticPool
from app.assets import lifecycle
from app.assets import manager as manager_module
from app.assets import scanner, seeder as seeder_module
from app.assets.database.models import Asset, AssetContent
from app.assets.database.queries.records import create_content, create_record
from app.assets.manager import AssetsEnabled
from app.assets.seeder import ScanStatus, asset_seeder
from app.database.models import Base
from app.assets.services.schemas import RegisteredAsset, UploadAssetView
class _ArgsStub:
enable_assets = True
enable_asset_hashing = False
class _OutputSeeder(Protocol):
def wait(self, timeout: float | None = None) -> bool: ...
def get_status(self) -> ScanStatus: ...
def shutdown(self, timeout: float = 5.0) -> bool: ...
@pytest.fixture
def enabled_manager() -> AssetsEnabled:
return AssetsEnabled(_ArgsStub())
@pytest.fixture
def asset_roots(
monkeypatch: pytest.MonkeyPatch, tmp_path: Path
) -> tuple[Path, Path, Path]:
output_dir = tmp_path / "output"
input_dir = tmp_path / "input"
temp_dir = tmp_path / "temp"
output_dir.mkdir()
input_dir.mkdir()
temp_dir.mkdir()
monkeypatch.setattr(folder_paths, "get_output_directory", lambda: str(output_dir))
monkeypatch.setattr(folder_paths, "get_input_directory", lambda: str(input_dir))
monkeypatch.setattr(folder_paths, "get_temp_directory", lambda: str(temp_dir))
return output_dir, input_dir, temp_dir
@pytest.fixture
def threaded_create_session(
monkeypatch: pytest.MonkeyPatch,
) -> Iterator[Callable[[], AbstractContextManager[Session]]]:
engine = create_engine(
"sqlite://",
connect_args={"check_same_thread": False},
poolclass=StaticPool,
)
Base.metadata.create_all(engine)
@contextmanager
def _create_session() -> Generator[SASession, None, None]:
with SASession(engine) as session:
yield session
monkeypatch.setattr(seeder_module, "create_session", _create_session)
monkeypatch.setattr(scanner, "create_session", _create_session)
monkeypatch.setattr("app.assets.services.ingest.create_session", _create_session)
yield _create_session
engine.dispose()
@pytest.fixture
def output_seeder(monkeypatch: pytest.MonkeyPatch) -> Iterator[_OutputSeeder]:
seeder = asset_seeder.__class__()
monkeypatch.setattr(manager_module, "asset_seeder", seeder)
yield seeder
_ = seeder.shutdown()
def test_queue_output_scan_registers_undeclared_output(
enabled_manager: AssetsEnabled,
asset_roots: tuple[Path, Path, Path],
threaded_create_session: Callable[[], AbstractContextManager[Session]],
output_seeder: _OutputSeeder,
) -> None:
output_dir, _, _ = asset_roots
output_path = output_dir / "undeclared.bin"
output_path.write_bytes(b"custom node output")
enabled_manager.queue_output_scan()
assert output_seeder.wait(timeout=5)
with threaded_create_session() as session:
rows = list(
session.scalars(
select(Asset)
.join(AssetContent, Asset.content_id == AssetContent.id)
.where(AssetContent.path == str(output_path.resolve()))
)
)
assert rows[0].job_id is None
assert len(rows) == 1
def test_queue_output_scan_does_not_duplicate_declared_output(
enabled_manager: AssetsEnabled,
asset_roots: tuple[Path, Path, Path],
threaded_create_session: Callable[[], AbstractContextManager[Session]],
output_seeder: _OutputSeeder,
) -> None:
output_dir, _, _ = asset_roots
output_path = output_dir / "declared.bin"
output_path.write_bytes(b"declared custom node output")
registered = enabled_manager.register_executed_output(
str(output_path), job_id="declared-job"
)
assert registered is not None
undeclared_path = output_dir / "alongside.bin"
undeclared_path.write_bytes(b"undeclared custom node output")
enabled_manager.queue_output_scan()
assert output_seeder.wait(timeout=5)
assert output_seeder.get_status().errors == []
with threaded_create_session() as session:
rows = list(
session.scalars(
select(Asset)
.join(AssetContent, Asset.content_id == AssetContent.id)
.where(AssetContent.path == str(output_path.resolve()))
)
)
undeclared_rows = list(
session.scalars(
select(Asset)
.join(AssetContent, Asset.content_id == AssetContent.id)
.where(AssetContent.path == str(undeclared_path.resolve()))
)
)
assert rows[0].job_id == "declared-job"
assert undeclared_rows[0].job_id is None
assert len(rows) == 1
assert len(undeclared_rows) == 1, (
"the undeclared sibling proves the walk ran, so the declared row's count of 1 is a real skip rather than a silently-failed scan"
)
def test_executed_and_cached_outputs_share_unhashed_content(
enabled_manager: AssetsEnabled,
asset_roots: tuple[Path, Path, Path],
mock_create_session: Callable[[], AbstractContextManager[Session]],
) -> None:
output_dir, _, _ = asset_roots
output_path = output_dir / "executed.png"
output_path.write_bytes(b"executed output")
executed = enabled_manager.register_executed_output(
str(output_path), job_id="executed-job"
)
assert isinstance(executed, RegisteredAsset)
assert executed.id
assert executed.content_id
assert executed.name
assert executed.job_id == "executed-job"
assert not hasattr(executed, "asset_hash")
assert {field.name for field in dataclasses.fields(executed)} == {
"id",
"content_id",
"job_id",
"name",
}
with mock_create_session() as session:
asset = session.get(Asset, executed.id)
content = session.get(AssetContent, executed.content_id)
assert asset is not None
assert asset.content_id == executed.content_id
assert content is not None
assert content.hash is None
cached = enabled_manager.register_cached_output(str(output_path), job_id="cached-job")
assert isinstance(cached, RegisteredAsset)
assert cached.id != executed.id
assert cached.content_id == executed.content_id
assert cached.job_id == "cached-job"
assert (
enabled_manager.register_cached_output(
str(output_dir / "unknown.png"), job_id="unknown-job"
)
is None
)
def test_register_upload_hashes_and_tags_fresh_input_file(
enabled_manager: AssetsEnabled,
asset_roots: tuple[Path, Path, Path],
mock_create_session: Callable[[], AbstractContextManager[Session]],
) -> None:
_, input_dir, _ = asset_roots
upload_path = input_dir / "pasted" / "upload.png"
upload_path.parent.mkdir()
upload_path.write_bytes(b"uploaded input")
view = enabled_manager.register_upload(
str(upload_path),
name=upload_path.name,
upload_type="input",
subfolder="pasted",
content_written=True,
)
assert isinstance(view, UploadAssetView)
assert isinstance(view.asset, RegisteredAsset)
assert view.asset_hash
assert "pasted" in view.tags
def test_startup_runs_against_memory_db_without_starting_a_scanner_thread(
enabled_manager: AssetsEnabled,
asset_roots: tuple[Path, Path, Path],
mock_create_session: Callable[[], AbstractContextManager[Session]],
monkeypatch: pytest.MonkeyPatch,
) -> None:
_, _, temp_dir = asset_roots
(temp_dir / "stale.tmp").write_bytes(b"stale")
seeder_start = MagicMock(return_value=False)
monkeypatch.setattr(lifecycle, "create_session", mock_create_session)
monkeypatch.setattr(lifecycle, "start_asset_seeder", seeder_start)
thread_count = threading.active_count()
enabled_manager.startup()
assert threading.active_count() == thread_count, (
"start_asset_seeder is mocked, so a new thread means a component other than the seeder spawned one"
)
assert not temp_dir.exists()
seeder_start.assert_called_once_with()
def test_ensure_scan_started_starts_the_lazy_object_info_scan(
enabled_manager: AssetsEnabled, monkeypatch: pytest.MonkeyPatch
) -> None:
seeder_start = MagicMock()
monkeypatch.setattr(asset_seeder, "start", seeder_start)
enabled_manager.ensure_scan_started()
seeder_start.assert_called_once_with(roots=("models", "input", "output"))
def test_shutdown_runs_lifecycle_cleanup_when_seeder_shutdown_times_out(
enabled_manager: AssetsEnabled,
asset_roots: tuple[Path, Path, Path],
mock_create_session: Callable[[], AbstractContextManager[Session]],
monkeypatch: pytest.MonkeyPatch,
) -> None:
_, _, temp_dir = asset_roots
stale_file = temp_dir / "stale.tmp"
stale_file.write_bytes(b"stale")
with mock_create_session() as session:
content = create_content(session, path=str(stale_file))
asset = create_record(session, content_id=content.id, name=stale_file.name)
session.commit()
asset_id = asset.id
content_id = content.id
calls = Mock()
seeder_shutdown = MagicMock(return_value=False)
run_shutdown = MagicMock(wraps=manager_module.run_shutdown)
calls.attach_mock(seeder_shutdown, "seeder_shutdown")
calls.attach_mock(run_shutdown, "run_shutdown")
monkeypatch.setattr(lifecycle, "can_create_session", lambda: True)
monkeypatch.setattr(lifecycle, "create_session", mock_create_session)
monkeypatch.setattr(asset_seeder, "shutdown", seeder_shutdown)
monkeypatch.setattr(manager_module, "run_shutdown", run_shutdown)
enabled_manager.shutdown()
assert calls.mock_calls == [call.seeder_shutdown(), call.run_shutdown()]
assert not temp_dir.exists()
with mock_create_session() as session:
assert session.get(Asset, asset_id) is None
assert session.get(AssetContent, content_id) is None

View File

@@ -0,0 +1,59 @@
import pytest
from app.assets.database.models import Asset, AssetTag, Tag
from app.assets.database.queries.records import (
create_content,
create_record,
fetch_record_tags,
)
from app.assets.services.asset_management import update_asset_metadata
from app.assets.services.tagging import apply_tags, remove_tags
def _create_record(session, path: str, tags: list[str] | None = None) -> Asset:
content = create_content(session, path)
record = create_record(session, content.id, "original", tags=tags)
session.commit()
return record
def test_rename_via_update_asset_metadata(session, mock_create_session):
record = _create_record(session, "/output/rename.png")
update_asset_metadata(record.id, name="renamed")
session.expire_all()
assert session.get(Asset, record.id).name == "renamed"
def test_apply_tags_adds_tags(session, mock_create_session):
record = _create_record(session, "/output/apply-tags.png")
apply_tags(record.id, ["foo", "bar"])
session.expire_all()
assert fetch_record_tags(session, record.id) == ["bar", "foo"]
def test_remove_tags_removes_non_automatic(session, mock_create_session):
record = _create_record(session, "/output/remove-tags.png", tags=["foo"])
session.add(Tag(name="automatic"))
session.add(AssetTag(asset_id=record.id, tag_name="automatic", origin="automatic"))
session.commit()
remove_tags(record.id, ["foo"])
session.expire_all()
assert fetch_record_tags(session, record.id) == ["automatic"]
def test_update_asset_metadata_unknown_preview_id_raises(session, mock_create_session):
record = _create_record(session, "/output/preview-validate.png")
with pytest.raises(ValueError):
update_asset_metadata(
record.id, preview_id="00000000-0000-0000-0000-000000000000"
)
session.expire_all()
assert session.get(Asset, record.id).preview_id is None

View File

@@ -0,0 +1,202 @@
from collections.abc import Callable
from contextlib import AbstractContextManager
import logging
from pathlib import Path
from unittest.mock import Mock, patch
import pytest
from aiohttp import web
from aiohttp.pytest_plugin import AiohttpClient
from sqlalchemy.orm import Session
import folder_paths
from app.assets import lifecycle
from app.assets.api import routes
from app.assets.database.models import Asset, AssetContent
from app.assets.database.queries.records import create_content, create_record
from app.assets import manager
from app.assets.manager import AssetsEnabled, NoAssets
from app.assets.mode import hashing_enabled
from app.assets.seeder import asset_seeder
from app.assets.services.hash_mode_state import read_stored_mode
class _Args:
enable_assets = False
def __init__(self, hashing: bool) -> None:
self.enable_asset_hashing = hashing
def _no_assets(*, hashing: bool = False) -> NoAssets:
return NoAssets(_Args(hashing))
@pytest.mark.asyncio
async def test_noassets_register_routes_returns_service_disabled_and_disables_seeder(
aiohttp_client: AiohttpClient, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(routes, "_ASSETS_ENABLED", False)
monkeypatch.setattr(asset_seeder, "_disabled", False)
app = web.Application()
_no_assets().register_routes(app, None)
client = await aiohttp_client(app)
response = await client.get("/api/assets")
assert response.status == 503
assert (await response.json())["error"]["code"] == "SERVICE_DISABLED"
assert asset_seeder.is_disabled()
def test_noassets_startup_applies_hash_mode_persists_state_and_cleans_temp_dir(
mock_create_session: Callable[[], AbstractContextManager[Session]],
session: Session,
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
temp_file = tmp_path / "startup-stale.bin"
temp_file.write_bytes(b"stale")
monkeypatch.setattr(lifecycle, "create_session", mock_create_session)
monkeypatch.setattr(folder_paths, "get_temp_directory", lambda: str(tmp_path))
_no_assets(hashing=True).startup()
assert hashing_enabled() is True
assert read_stored_mode(session) == "on"
assert not temp_file.exists()
def test_noassets_shutdown_wipes_temp_rows_and_files(
mock_create_session: Callable[[], AbstractContextManager[Session]],
session: Session,
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
manager = _no_assets()
monkeypatch.setattr(lifecycle, "create_session", mock_create_session)
monkeypatch.setattr(lifecycle, "can_create_session", lambda: True)
monkeypatch.setattr(folder_paths, "get_temp_directory", lambda: str(tmp_path))
manager.startup()
tmp_path.mkdir()
temp_file = tmp_path / "shutdown-stale.bin"
temp_file.write_bytes(b"stale")
content = create_content(session, str(temp_file))
record = create_record(session, content.id, temp_file.name)
session.commit()
record_id = record.id
content_id = content.id
manager.shutdown()
session.expire_all()
assert session.get(Asset, record_id) is None
assert session.get(AssetContent, content_id) is None
assert not temp_file.exists()
def test_noassets_shutdown_runs_cleanup_after_unsuccessful_seeder_shutdown(
mock_create_session: Callable[[], AbstractContextManager[Session]],
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
temp_file = tmp_path / "seeder-timeout.bin"
temp_file.write_bytes(b"stale")
calls: list[str] = []
monkeypatch.setattr(lifecycle, "create_session", mock_create_session)
monkeypatch.setattr(lifecycle, "can_create_session", lambda: True)
monkeypatch.setattr(folder_paths, "get_temp_directory", lambda: str(tmp_path))
def seeder_shutdown() -> bool:
calls.append("seeder")
return False
def cleanup() -> None:
calls.append("cleanup")
lifecycle.run_shutdown()
with (
patch.object(asset_seeder, "shutdown", side_effect=seeder_shutdown),
patch("app.assets.manager.run_shutdown", side_effect=cleanup) as cleanup_spy,
):
_no_assets().shutdown()
cleanup_spy.assert_called_once_with()
assert calls == ["seeder", "cleanup"]
assert not temp_file.exists()
def test_noassets_shutdown_without_database_sweeps_files_without_session_access(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
temp_file = tmp_path / "no-database.bin"
temp_file.write_bytes(b"stale")
create_session = Mock()
monkeypatch.setattr(lifecycle, "can_create_session", lambda: False)
monkeypatch.setattr(lifecycle, "create_session", create_session)
monkeypatch.setattr(folder_paths, "get_temp_directory", lambda: str(tmp_path))
_no_assets().shutdown()
assert not temp_file.exists()
create_session.assert_not_called()
def test_noassets_callbacks_are_noops_without_seeder_side_effects() -> None:
manager = _no_assets()
with (
patch.object(asset_seeder, "start") as start,
patch.object(asset_seeder, "pause") as pause,
patch.object(asset_seeder, "enqueue_scan") as enqueue_scan,
patch.object(asset_seeder, "resume") as resume,
patch.object(asset_seeder, "set_event_sink") as set_event_sink,
):
assert manager.ensure_scan_started() is None
assert manager.pause_background_scan() is None
assert manager.queue_output_scan() is None
assert manager.resume_background_scan() is None
assert manager.set_event_sink(Mock()) is None
start.assert_not_called()
pause.assert_not_called()
enqueue_scan.assert_not_called()
resume.assert_not_called()
set_event_sink.assert_not_called()
def test_noassets_registration_methods_return_none() -> None:
manager = _no_assets()
assert manager.register_upload(
"/tmp/output.png", "output.png", "output", "", content_written=True
) is None
assert manager.register_executed_output("/tmp/output.png", "job-id") is None
assert manager.register_cached_output("/tmp/output.png", "job-id") is None
def test_noassets_is_disabled() -> None:
assert _no_assets().enabled is False
def test_default_asset_manager_disables_assets_when_dependencies_are_unavailable(
monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
) -> None:
monkeypatch.setattr(manager.args, "enable_assets", True)
monkeypatch.setattr(manager, "dependencies_available", lambda: False)
with caplog.at_level(logging.WARNING):
asset_manager = manager.default_asset_manager()
assert isinstance(asset_manager, NoAssets)
assert "asset endpoints will answer 503" in caplog.text
assert "requirements.txt" in caplog.text
def test_default_asset_manager_enables_assets_when_dependencies_are_available(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(manager.args, "enable_assets", True)
monkeypatch.setattr(manager, "dependencies_available", lambda: True)
assert isinstance(manager.default_asset_manager(), AssetsEnabled)

View File

@@ -0,0 +1,198 @@
import os
from pathlib import Path
from unittest.mock import patch
import pytest
from sqlalchemy import select
from app.assets.database.models import Asset, AssetContent, AssetTag
from app.assets.database.queries.records import create_content, create_record
from app.assets.helpers import to_stored_hash
from app.assets.scanner import SeedAssetSpec, clear_pending_verifications, seed_asset_specs
from app.assets.services import hash_mode_state
from app.assets.services.hash_mode_state import (
clear_transition_queue,
drain_transition_queue,
enqueue_transition_work,
record_transition_intent,
write_stored_mode,
)
from app.assets.services.snapshot_hash import snapshot_hash
@pytest.fixture(autouse=True)
def _clear_queues():
clear_transition_queue()
clear_pending_verifications()
yield
clear_transition_queue()
clear_pending_verifications()
def _stored_hash(path: Path) -> str:
snapshot = snapshot_hash(str(path))
assert snapshot is not None
digest, _ = snapshot
return to_stored_hash(digest)
def _spec(path: Path) -> SeedAssetSpec:
stat = path.stat()
return {
"abs_path": str(path),
"size_bytes": stat.st_size,
"mtime_ns": stat.st_mtime_ns,
"info_name": path.name,
"tags": ["input"],
"fname": path.name,
"metadata": None,
"mime_type": None,
"job_id": None,
}
def test_deleted_null_hash_row_recovers_via_scanner_after_restore(
session, temp_dir, monkeypatch
):
path = temp_dir / "recoverable.bin"
monkeypatch.setattr("folder_paths.get_input_directory", lambda: str(temp_dir))
original_bytes = b"the exact bytes that come back"
path.write_bytes(original_bytes)
stat = path.stat()
content = create_content(session, str(path), size_bytes=stat.st_size, mtime_ns=stat.st_mtime_ns)
content_id = content.id
record = create_record(session, content_id, "recoverable.bin", tags=["input"])
record_id = record.id
session.commit()
write_stored_mode(session, "off")
monkeypatch.setattr(hash_mode_state._mode, "hashing_enabled", lambda: True)
path.unlink()
transition = record_transition_intent(session)
enqueue_transition_work(session, transition)
drain_transition_queue(session)
session.commit()
assert session.get(AssetContent, content_id).is_missing is True, (
"precondition: A1's drain marked the row missing on delete"
)
path.write_bytes(original_bytes)
os.utime(path, ns=(stat.st_atime_ns, stat.st_mtime_ns))
assert path.stat().st_mtime_ns == stat.st_mtime_ns, "setup: mtime must round-trip exactly"
with patch("app.assets.scanner.mode.hashing_enabled", return_value=True):
created = seed_asset_specs(session, [_spec(path)])
session.commit()
assert created == 0, "the original row must recover — no fresh content row minted"
recovered = session.get(AssetContent, content_id)
assert recovered.is_missing is False
assert recovered.hash == _stored_hash(path)
assert len(session.scalars(select(AssetContent)).all()) == 1, "no duplicate content row"
assert len(session.scalars(select(Asset)).all()) == 1, "no duplicate record"
assert session.get(Asset, record_id) is not None, "the original record's identity survives"
assert session.get(AssetTag, {"asset_id": record_id, "tag_name": "missing"}) is None
def test_different_bytes_restored_at_same_path_does_not_recover_old_row(
session, temp_dir, monkeypatch
):
path = temp_dir / "replaced.bin"
monkeypatch.setattr("folder_paths.get_input_directory", lambda: str(temp_dir))
path.write_bytes(b"original bytes")
stat = path.stat()
content = create_content(session, str(path), size_bytes=stat.st_size, mtime_ns=stat.st_mtime_ns)
content_id = content.id
create_record(session, content_id, "replaced.bin")
session.commit()
write_stored_mode(session, "off")
monkeypatch.setattr(hash_mode_state._mode, "hashing_enabled", lambda: True)
path.unlink()
transition = record_transition_intent(session)
enqueue_transition_work(session, transition)
drain_transition_queue(session)
session.commit()
assert session.get(AssetContent, content_id).is_missing is True
path.write_bytes(b"a completely different, much longer payload than the original")
with patch("app.assets.scanner.mode.hashing_enabled", return_value=True):
created = seed_asset_specs(session, [_spec(path)])
session.commit()
assert created == 1, "a genuinely different file must take the normal new-content path"
assert session.get(AssetContent, content_id).is_missing is True, (
"the old row must stay missing — recovering it here would hand the wrong record's "
"identity, tags, and metadata to unrelated bytes"
)
assert len(session.scalars(select(AssetContent)).all()) == 2
def test_same_size_different_mtime_restored_at_same_path_does_not_recover_old_row(
session, temp_dir, monkeypatch
):
path = temp_dir / "retimed.bin"
monkeypatch.setattr("folder_paths.get_input_directory", lambda: str(temp_dir))
original_bytes = b"identical length, different moment in time"
path.write_bytes(original_bytes)
stat = path.stat()
content = create_content(session, str(path), size_bytes=stat.st_size, mtime_ns=stat.st_mtime_ns)
content_id = content.id
create_record(session, content_id, "retimed.bin")
session.commit()
write_stored_mode(session, "off")
monkeypatch.setattr(hash_mode_state._mode, "hashing_enabled", lambda: True)
path.unlink()
transition = record_transition_intent(session)
enqueue_transition_work(session, transition)
drain_transition_queue(session)
session.commit()
assert session.get(AssetContent, content_id).is_missing is True
path.write_bytes(original_bytes)
shifted_ns = stat.st_mtime_ns + 5_000_000_000
os.utime(path, ns=(stat.st_atime_ns, shifted_ns))
assert path.stat().st_mtime_ns != stat.st_mtime_ns, "setup: mtime must actually differ"
assert path.stat().st_size == stat.st_size, "setup: size must match so only mtime disambiguates"
with patch("app.assets.scanner.mode.hashing_enabled", return_value=True):
created = seed_asset_specs(session, [_spec(path)])
session.commit()
assert created == 1, "a same-size-but-different-mtime restore must take the new-content path"
assert session.get(AssetContent, content_id).is_missing is True, (
"a matching size with a mismatched mtime is not proof the old row's bytes are back — "
"recovering here would hand the wrong record's identity to different bytes"
)
assert len(session.scalars(select(AssetContent)).all()) == 2
def test_two_missing_null_hash_candidates_at_same_path_do_not_recover(
session, temp_dir, monkeypatch
):
path = temp_dir / "ambiguous_null.bin"
monkeypatch.setattr("folder_paths.get_input_directory", lambda: str(temp_dir))
path.write_bytes(b"bytes shared by two missing generations")
stat = path.stat()
first = AssetContent(
path=str(path), hash=None, is_missing=True,
size_bytes=stat.st_size, mtime_ns=stat.st_mtime_ns,
)
second = AssetContent(
path=str(path), hash=None, is_missing=True,
size_bytes=stat.st_size, mtime_ns=stat.st_mtime_ns,
)
session.add_all([first, second])
session.flush()
first_id, second_id = first.id, second.id
session.commit()
with patch("app.assets.scanner.mode.hashing_enabled", return_value=True):
created = seed_asset_specs(session, [_spec(path)])
session.commit()
assert created == 1, "ambiguous candidates must fall through to the normal new-content path"
assert session.get(AssetContent, first_id).is_missing is True
assert session.get(AssetContent, second_id).is_missing is True

View File

@@ -0,0 +1,143 @@
import os
from pathlib import Path
from unittest.mock import patch
import folder_paths
from sqlalchemy import select
from app.assets.database.models import Asset, AssetContent
from app.assets.manager import AssetsEnabled, NoAssets
from comfy_execution.asset_enrichment import (
register_cached_outputs,
register_executed_outputs,
)
class _ArgsStub:
def __init__(self, enable_assets: bool = True) -> None:
self.enable_assets = enable_assets
self.enable_asset_hashing = False
def _write_output_file(name: str, data: bytes) -> Path:
path = Path(folder_paths.get_output_directory()) / name
path.parent.mkdir(parents=True, exist_ok=True)
path.write_bytes(data)
return path
def _output_ui(name: str) -> dict:
return {"images": [{"filename": name, "subfolder": "", "type": "output"}]}
def _wrapper(name: str, node_id: str = "1") -> dict:
return {"meta": {"node_id": node_id}, "output": _output_ui(name)}
def test_executed_adapter_registers_new_output(mock_create_session):
path = _write_output_file("adapter-executed-new.png", b"exec")
try:
output_ui = _output_ui(path.name)
enriched = register_executed_outputs(
output_ui, "exec-job", AssetsEnabled(_ArgsStub())
)
new_id = enriched["images"][0]["id"]
assert "id" not in output_ui["images"][0]
with mock_create_session() as session:
record = session.get(Asset, new_id)
content = session.get(AssetContent, record.content_id)
assert record.job_id == "exec-job"
assert content.path == os.path.abspath(path)
assert content.is_missing is False
finally:
path.unlink(missing_ok=True)
def test_executed_adapter_over_existing_path_marks_old_missing(mock_create_session):
path = _write_output_file("adapter-executed-replace.png", b"original")
try:
manager = AssetsEnabled(_ArgsStub())
first = register_executed_outputs(_output_ui(path.name), "job-1", manager)
old_id = first["images"][0]["id"]
path.write_bytes(b"replacement")
second = register_executed_outputs(_output_ui(path.name), "job-2", manager)
new_id = second["images"][0]["id"]
assert new_id != old_id
with mock_create_session() as session:
old_content = session.get(AssetContent, session.get(Asset, old_id).content_id)
new_content = session.get(AssetContent, session.get(Asset, new_id).content_id)
assert old_content.is_missing is True
assert new_content.is_missing is False
contents = list(
session.scalars(
select(AssetContent).where(AssetContent.path == os.path.abspath(path))
)
)
assert len(contents) == 2
finally:
path.unlink(missing_ok=True)
def test_executed_adapter_disabled_registers_nothing(mock_create_session):
path = _write_output_file("adapter-executed-disabled.png", b"noop")
try:
output_ui = _output_ui(path.name)
enriched = register_executed_outputs(output_ui, "job", NoAssets(_ArgsStub(False)))
assert "id" not in enriched["images"][0]
with mock_create_session() as session:
contents = list(
session.scalars(
select(AssetContent).where(AssetContent.path == os.path.abspath(path))
)
)
assert contents == []
finally:
path.unlink(missing_ok=True)
def test_executed_adapter_registration_failure_never_raises(mock_create_session):
path = _write_output_file("adapter-executed-error.png", b"exec")
try:
output_ui = _output_ui(path.name)
with patch(
"app.assets.manager.ingest_register_executed_output",
side_effect=RuntimeError("boom"),
):
enriched = register_executed_outputs(
output_ui, "job", AssetsEnabled(_ArgsStub())
)
assert "id" not in enriched["images"][0]
finally:
path.unlink(missing_ok=True)
def test_cached_adapter_without_live_content_is_nonevent(mock_create_session):
path = _write_output_file("adapter-cached-missing.png", b"new content")
try:
wrapper = _wrapper(path.name)
enriched = register_cached_outputs(
wrapper, "cached-job", AssetsEnabled(_ArgsStub())
)
assert "id" not in enriched["output"]["images"][0]
with mock_create_session() as session:
contents = list(
session.scalars(
select(AssetContent).where(AssetContent.path == os.path.abspath(path))
)
)
records = list(
session.scalars(select(Asset).where(Asset.job_id == "cached-job"))
)
assert contents == []
assert records == []
finally:
path.unlink(missing_ok=True)

View File

@@ -0,0 +1,15 @@
import subprocess
from pathlib import Path
def test_no_owner_id_in_asset_modules():
repository = Path(__file__).resolve().parents[3]
result = subprocess.run(
["grep", "-rn", "--exclude=*.pyc", "owner_id", "app/assets/"],
cwd=repository,
check=False,
capture_output=True,
text=True,
)
assert result.stdout == ""

View File

@@ -0,0 +1,64 @@
from sqlalchemy import event
from app.assets.database.queries.records import (
create_content,
create_record,
mark_content_missing,
)
from app.assets.services.asset_management import get_preview_file_paths
def test_preview_paths_resolve_preview_record_content(session, mock_create_session) -> None:
preview_content = create_content(session, "/output/preview.png")
preview = create_record(session, preview_content.id, "preview.png")
record_content = create_content(session, "/output/record.png")
record = create_record(session, record_content.id, "record.png")
record.preview_id = preview.id
session.commit()
paths = get_preview_file_paths([preview.id])
assert paths == {preview.id: preview_content.path}
def test_preview_paths_exclude_missing_preview_content(session, mock_create_session) -> None:
preview_content = create_content(session, "/output/missing-preview.png")
preview = create_record(session, preview_content.id, "missing-preview.png")
record_content = create_content(session, "/output/record.png")
record = create_record(session, record_content.id, "record.png")
record.preview_id = preview.id
mark_content_missing(session, preview_content.id)
session.commit()
paths = get_preview_file_paths([preview.id])
assert paths == {}
def test_preview_paths_resolve_a_page_in_one_query(session, mock_create_session, db_engine) -> None:
preview_ids: list[str] = []
expected_paths: dict[str, str] = {}
for index in range(3):
preview_path = f"/output/preview-{index}.png"
preview_content = create_content(session, preview_path)
preview = create_record(session, preview_content.id, f"preview-{index}.png")
record_content = create_content(session, f"/output/record-{index}.png")
record = create_record(session, record_content.id, f"record-{index}.png")
record.preview_id = preview.id
preview_ids.append(preview.id)
expected_paths[preview.id] = preview_content.path
session.commit()
statements: list[str] = []
def count_statements(_, __, statement, ___, ____, _____) -> None:
statements.append(statement)
event.listen(db_engine, "before_cursor_execute", count_statements)
try:
paths = get_preview_file_paths(preview_ids)
finally:
event.remove(db_engine, "before_cursor_execute", count_statements)
assert paths == expected_paths
assert len(statements) == 1

View File

@@ -0,0 +1,168 @@
from pathlib import Path
from unittest.mock import patch
import pytest
from sqlalchemy import select
from app.assets.database.models import Asset, AssetContent, AssetTag, Tag
from app.assets.helpers import to_stored_hash
from app.assets.scanner import (
SeedAssetSpec,
clear_pending_verifications,
pending_recovery_count,
seed_asset_specs,
)
from app.assets.scanner_changes import recover_missing_content
from app.assets.services.snapshot_hash import snapshot_hash
@pytest.fixture(autouse=True)
def _clear_queues():
clear_pending_verifications()
yield
clear_pending_verifications()
def _missing_content(session, path: Path, hash_value: str) -> tuple[AssetContent, Asset]:
content = AssetContent(path=str(path), hash=hash_value, is_missing=True)
session.add(content)
session.flush()
record = Asset(content_id=content.id, name=path.name)
session.add(record)
session.flush()
if session.get(Tag, "missing") is None:
session.add(Tag(name="missing"))
session.add(AssetTag(asset_id=record.id, tag_name="missing", origin="automatic"))
session.commit()
return content, record
def _spec(path: Path) -> SeedAssetSpec:
stat = path.stat()
return {
"abs_path": str(path),
"size_bytes": stat.st_size,
"mtime_ns": stat.st_mtime_ns,
"info_name": path.name,
"tags": ["input"],
"fname": path.name,
"metadata": None,
"mime_type": None,
"job_id": None,
}
def _stored_hash(path: Path) -> str:
snapshot = snapshot_hash(str(path))
assert snapshot is not None
digest, _ = snapshot
return to_stored_hash(digest)
def test_single_hash_match_recovers(session, temp_dir: Path):
path = temp_dir / "restored.bin"
path.write_bytes(b"restored bytes")
content, record = _missing_content(session, path, _stored_hash(path))
with patch("app.assets.scanner.mode.hashing_enabled", return_value=True):
created = seed_asset_specs(session, [_spec(path)])
session.commit()
assert created == 0
assert session.get(AssetContent, content.id).is_missing is False
assert session.get(AssetTag, {"asset_id": record.id, "tag_name": "missing"}) is None
def test_ambiguous_hash_match_recovers_nothing(session, temp_dir: Path):
path = temp_dir / "ambiguous.bin"
path.write_bytes(b"same bytes")
digest = _stored_hash(path)
first, _ = _missing_content(session, path, digest)
second, _ = _missing_content(session, path, digest)
with patch("app.assets.scanner.mode.hashing_enabled", return_value=True):
created = seed_asset_specs(session, [_spec(path)])
session.commit()
assert created == 1
assert session.get(AssetContent, first.id).is_missing is True
assert session.get(AssetContent, second.id).is_missing is True
assert len(session.scalars(select(AssetContent)).all()) == 3
def test_no_hash_match_creates_fresh_rows(session, temp_dir: Path):
path = temp_dir / "different.bin"
path.write_bytes(b"current bytes")
missing, _ = _missing_content(session, path, "old")
with patch("app.assets.scanner.mode.hashing_enabled", return_value=True):
created = seed_asset_specs(session, [_spec(path)])
session.commit()
assert created == 1
assert session.get(AssetContent, missing.id).is_missing is True
assert len(session.scalars(select(Asset)).all()) == 2
def test_off_mode_no_recovery(session, temp_dir: Path):
path = temp_dir / "off.bin"
path.write_bytes(b"bytes")
missing, _ = _missing_content(session, path, _stored_hash(path))
with (
patch("app.assets.scanner.mode.hashing_enabled", return_value=False),
patch("app.assets.scanner_changes.snapshot_hash") as hash_mock,
):
created = seed_asset_specs(session, [_spec(path)])
session.commit()
hash_mock.assert_not_called()
assert created == 1
assert session.get(AssetContent, missing.id).is_missing is True
def test_unstable_hash_requeues(session, temp_dir: Path):
path = temp_dir / "unstable.bin"
path.write_bytes(b"bytes")
missing, _ = _missing_content(session, path, "old")
with (
patch("app.assets.scanner.mode.hashing_enabled", return_value=True),
patch("app.assets.scanner_changes.snapshot_hash", return_value=None),
):
created = seed_asset_specs(session, [_spec(path)])
session.commit()
assert created == 0
assert pending_recovery_count() == 1
assert session.get(AssetContent, missing.id).is_missing is True
def test_recovery_skips_a_path_a_live_row_already_occupies(session, temp_dir: Path):
path = temp_dir / "contested.bin"
path.write_bytes(b"bytes restored under a path someone else already owns")
digest = _stored_hash(path)
stat_result = path.stat()
missing_a, _ = _missing_content(session, path, digest)
live_b = AssetContent(
path=str(path),
hash=digest,
is_missing=False,
size_bytes=stat_result.st_size,
mtime_ns=stat_result.st_mtime_ns,
)
session.add(live_b)
session.commit()
live_b_id = live_b.id
result = recover_missing_content(
session, str(path), stat_result, hashing_is_enabled=True
)
session.commit()
assert result == "no_match", (
"the file at an occupied path belongs to the live row's story; change detection "
"owns it, and recovering onto it would put two live rows on one path"
)
assert session.get(AssetContent, missing_a.id).is_missing is True
assert session.get(AssetContent, live_b_id).is_missing is False

View File

@@ -0,0 +1,273 @@
import os
from datetime import datetime
import folder_paths
from sqlalchemy import select
from app.assets.database.models import Asset, AssetContent
from app.assets.database.queries.records import create_content, create_record
from app.assets.services.ingest import (
register_cached_output,
register_executed_output,
register_file_in_place,
)
def _output_path(name: str) -> str:
output_dir = folder_paths.get_output_directory()
os.makedirs(output_dir, exist_ok=True)
return os.path.join(output_dir, name)
def test_executed_creates_metadata_and_null_hash(mock_create_session):
body = b"pixels"
path = _output_path("primitives_executed_meta.png")
with open(path, "wb") as fh:
fh.write(body)
try:
record = register_executed_output(path, job_id="job-exec")
assert record is not None
with mock_create_session() as session:
content = session.scalar(
select(AssetContent).where(
AssetContent.path == os.path.abspath(path)
)
)
asset = session.get(Asset, record.id)
assert content.hash is None
assert content.is_missing is False
assert asset.system_metadata is not None
assert asset.system_metadata.get("format") == "png"
assert asset.system_metadata.get("content_length") == len(body)
assert "filename" in asset.system_metadata
finally:
if os.path.exists(path):
os.unlink(path)
def test_executed_over_live_path_marks_old_missing(mock_create_session):
path = _output_path("primitives_executed_overwrite.png")
try:
with open(path, "wb") as fh:
fh.write(b"v1")
first = register_executed_output(path, job_id="job1")
with open(path, "wb") as fh:
fh.write(b"v2")
second = register_executed_output(path, job_id="job2")
assert first is not None
assert second is not None
assert second.id != first.id
with mock_create_session() as session:
old_content = session.get(AssetContent, first.content_id)
old_record = session.get(Asset, first.id)
new_record = session.get(Asset, second.id)
assert old_content.is_missing is True
assert old_record.job_id == "job1"
assert new_record.job_id == "job2"
assert new_record.content_id != first.content_id
finally:
if os.path.exists(path):
os.unlink(path)
def test_cached_copies_sibling_metadata(mock_create_session):
sibling_metadata = {"provenance": "sibling", "width": 4, "kind": "image"}
path = _output_path("primitives_cached_copy.png")
with open(path, "wb") as fh:
fh.write(b"pixels")
try:
with mock_create_session() as session:
content = create_content(session, os.path.abspath(path), None, 6, 0)
create_record(
session,
content.id,
"sibling",
system_metadata=dict(sibling_metadata),
)
session.commit()
content_id = content.id
cached = register_cached_output(path, job_id="delivery-job")
assert cached is not None
assert cached.content_id == content_id
with mock_create_session() as session:
record = session.get(Asset, cached.id)
assert record.system_metadata == sibling_metadata
finally:
if os.path.exists(path):
os.unlink(path)
def test_cached_earliest_sibling_wins(mock_create_session):
earliest = datetime(2020, 1, 1, 0, 0, 0)
later = datetime(2020, 1, 2, 0, 0, 0)
winner_metadata = {"winner": "earliest-low-id"}
path = _output_path("primitives_cached_earliest.png")
with open(path, "wb") as fh:
fh.write(b"pixels")
try:
with mock_create_session() as session:
content = create_content(session, os.path.abspath(path), None, 6, 0)
session.add_all(
[
Asset(
id="00000000-0000-0000-0000-000000000001",
content_id=content.id,
name="earliest-low-id",
created_at=earliest,
system_metadata=dict(winner_metadata),
),
Asset(
id="00000000-0000-0000-0000-000000000002",
content_id=content.id,
name="earliest-high-id",
created_at=earliest,
system_metadata={"winner": "earliest-high-id"},
),
Asset(
id="00000000-0000-0000-0000-000000000000",
content_id=content.id,
name="later-lowest-id",
created_at=later,
system_metadata={"winner": "later-lowest-id"},
),
]
)
session.commit()
content_id = content.id
cached = register_cached_output(path, job_id="delivery-job")
assert cached is not None
assert cached.content_id == content_id
with mock_create_session() as session:
record = session.get(Asset, cached.id)
assert record.system_metadata == winner_metadata
finally:
if os.path.exists(path):
os.unlink(path)
def test_cached_orphan_content_extracts_fresh(mock_create_session):
body = b"pixels"
path = _output_path("primitives_cached_orphan.png")
with open(path, "wb") as fh:
fh.write(body)
try:
with mock_create_session() as session:
content = create_content(
session, os.path.abspath(path), None, len(body), 0
)
session.commit()
content_id = content.id
cached = register_cached_output(path, job_id="delivery-job")
assert cached is not None
assert cached.content_id == content_id
with mock_create_session() as session:
record = session.get(Asset, cached.id)
assert record.system_metadata is not None
assert record.system_metadata.get("format") == "png"
assert record.system_metadata.get("content_length") == len(body)
finally:
if os.path.exists(path):
os.unlink(path)
def test_cached_missing_content_is_nonevent(mock_create_session):
path = _output_path("primitives_cached_missing.png")
with open(path, "wb") as fh:
fh.write(b"pixels")
try:
cached = register_cached_output(path, job_id="delivery-job")
assert cached is None
with mock_create_session() as session:
contents = list(
session.scalars(
select(AssetContent).where(
AssetContent.path == os.path.abspath(path)
)
)
)
records = list(
session.scalars(
select(Asset).where(Asset.job_id == "delivery-job")
)
)
assert contents == []
assert records == []
finally:
if os.path.exists(path):
os.unlink(path)
def test_registration_failure_never_raises_and_leaves_no_rows(
mock_create_session, monkeypatch
):
def _boom(*args, **kwargs):
raise RuntimeError("simulated create_record failure")
monkeypatch.setattr(
"app.assets.services.ingest.create_record", _boom
)
path = _output_path("primitives_reg_fail.png")
with open(path, "wb") as fh:
fh.write(b"pixels")
try:
result = register_executed_output(path, job_id="job-boom")
assert result is None
with mock_create_session() as session:
contents = list(
session.scalars(
select(AssetContent).where(
AssetContent.path == os.path.abspath(path)
)
)
)
records = list(
session.scalars(
select(Asset).where(Asset.job_id == "job-boom")
)
)
assert contents == []
assert records == []
finally:
if os.path.exists(path):
os.unlink(path)
def test_upload_record_has_metadata_and_hash(mock_create_session):
body = b"pixels-content"
path = _output_path("primitives_upload_in_place.png")
with open(path, "wb") as fh:
fh.write(body)
try:
result = register_file_in_place(
path, name="pic.png", tags=["demo"], mime_type="image/png"
)
assert result.asset.hash is not None
assert result.ref.system_metadata is not None
assert result.ref.system_metadata.get("format") == "png"
assert result.ref.system_metadata.get("content_length") == len(body)
finally:
if os.path.exists(path):
os.unlink(path)

View File

@@ -0,0 +1,43 @@
from pathlib import Path
from unittest.mock import patch
from sqlalchemy import select
from app.assets.database.models import AssetContent
from app.assets.scanner import build_asset_specs, seed_asset_specs, sync_prefixes_with_filesystem
def _scan(session, root: Path) -> int:
paths = [str(path) for path in root.iterdir()]
specs, _, _ = build_asset_specs(paths, set(), enable_metadata_extraction=False)
return seed_asset_specs(session, specs)
def test_e2e_scan_seed_detect_prune(session, temp_dir: Path):
root = temp_dir / "input"
root.mkdir()
removed = root / "removed.bin"
edited = root / "edited.bin"
removed.write_bytes(b"removed")
edited.write_bytes(b"old")
with patch("folder_paths.get_input_directory", return_value=str(root)):
assert _scan(session, root) == 2
removed.unlink()
edited.write_bytes(b"replacement")
(root / "partial.part").write_bytes(b"partial")
with patch("app.assets.scanner.mode.hashing_enabled", return_value=False):
sync_prefixes_with_filesystem(session, [str(root)])
_scan(session, root)
session.commit()
contents = list(session.scalars(select(AssetContent)))
assert len(contents) == 3
assert len([content for content in contents if content.is_missing]) == 2
def test_second_scan_idempotent(session, temp_dir: Path):
root = temp_dir / "input"
root.mkdir()
(root / "stable.bin").write_bytes(b"stable")
with patch("folder_paths.get_input_directory", return_value=str(root)):
assert _scan(session, root) == 1
assert _scan(session, root) == 0

View File

@@ -0,0 +1,232 @@
import os
from collections.abc import Mapping
from dataclasses import dataclass
from pathlib import Path
from unittest.mock import patch
from sqlalchemy import select
from sqlalchemy.orm import Session
from app.assets.database.models import Asset, AssetContent, AssetTag
from app.assets.helpers import to_stored_hash
from app.assets.scanner import (
build_asset_specs,
enrich_asset,
mark_contents_missing_outside_prefixes,
seed_asset_specs,
sync_prefixes_with_filesystem,
)
from app.assets.services.snapshot_hash import snapshot_hash
@dataclass(frozen=True, slots=True)
class _ExtractedMetadata:
content_type: str | None
system_metadata: Mapping[str, int | str]
def to_user_metadata(self) -> dict[str, int | str]:
return dict(self.system_metadata)
def _build_seed_specs(root: Path) -> list:
return build_asset_specs(
[str(path) for path in sorted(root.iterdir())],
existing_paths=set(),
enable_metadata_extraction=False,
)[0]
def _create_enrichment_target(
session: Session,
path: Path,
system_metadata: dict[str, int | str] | None = None,
) -> tuple[AssetContent, Asset]:
content = AssetContent(
path=str(path),
hash=None,
size_bytes=path.stat().st_size,
mtime_ns=path.stat().st_mtime_ns,
)
session.add(content)
session.flush()
record = Asset(
content_id=content.id,
name=path.name,
system_metadata=system_metadata,
)
session.add(record)
session.commit()
return content, record
def test_enrichment_retains_absent_system_metadata_keys(session: Session, temp_dir: Path):
path = temp_dir / "metadata.bin"
path.write_bytes(b"metadata")
content, record = _create_enrichment_target(session, path)
with patch(
"app.assets.scanner.extract_file_metadata",
side_effect=[
_ExtractedMetadata(None, {"a": 1, "b": 2}),
_ExtractedMetadata(None, {"b": 3}),
],
):
enrich_asset(session, str(path), content.id, record.id)
enrich_asset(session, str(path), content.id, record.id)
assert record.system_metadata == {"a": 1, "b": 3}
def test_enrichment_retains_dimensions_when_image_extraction_degrades(
session: Session, temp_dir: Path
):
path = temp_dir / "image.png"
path.write_bytes(b"not a complete image")
content, record = _create_enrichment_target(
session,
path,
{"kind": "image", "width": 64, "height": 64},
)
with (
patch(
"app.assets.scanner.extract_file_metadata",
return_value=_ExtractedMetadata("image/png", {"filename": "image.png"}),
),
patch("app.assets.scanner.extract_image_dimensions", return_value=None),
):
enrich_asset(session, str(path), content.id, record.id)
assert record.system_metadata == {
"filename": "image.png",
"kind": "image",
"width": 64,
"height": 64,
}
def test_enrichment_overrides_content_length_with_zero(
session: Session, temp_dir: Path
):
path = temp_dir / "empty.bin"
path.write_bytes(b"")
content, record = _create_enrichment_target(
session,
path,
{"content_length": 1024},
)
with patch(
"app.assets.scanner.extract_file_metadata",
return_value=_ExtractedMetadata(None, {"content_length": 0}),
):
enrich_asset(session, str(path), content.id, record.id)
assert record.system_metadata == {"content_length": 0}
def test_seed_creates_content_and_record(session, temp_dir: Path):
input_root = temp_dir / "input"
input_root.mkdir()
(input_root / "first.png").write_bytes(b"first")
(input_root / "second.png").write_bytes(b"second")
with patch("folder_paths.get_input_directory", return_value=str(input_root)):
created = seed_asset_specs(session, _build_seed_specs(input_root))
session.commit()
contents = list(session.scalars(select(AssetContent).order_by(AssetContent.path)))
records = list(session.scalars(select(Asset).order_by(Asset.name)))
assert created == 2
assert len(contents) == 2
assert len(records) == 2
assert [record.loader_path for record in records] == ["first.png", "second.png"]
assert [link.tag_name for link in session.scalars(select(AssetTag).order_by(AssetTag.asset_id))] == [
"input",
"input",
]
def test_prune_marks_missing_not_deletes(session, temp_dir: Path):
input_root = temp_dir / "input"
input_root.mkdir()
file_path = input_root / "removed-from-registry.png"
file_path.write_bytes(b"content")
with patch("folder_paths.get_input_directory", return_value=str(input_root)):
seed_asset_specs(session, _build_seed_specs(input_root))
session.commit()
marked = mark_contents_missing_outside_prefixes(session, prefixes=[])
session.commit()
content = session.scalar(select(AssetContent))
record = session.scalar(select(Asset))
missing_tag = session.get(AssetTag, {"asset_id": record.id, "tag_name": "missing"})
assert marked == 1
assert content is not None and content.is_missing is True
assert record is not None
assert missing_tag is not None and missing_tag.origin == "automatic"
def test_unhashed_missing_content_gets_tagged(session, temp_dir: Path):
missing_path = os.path.abspath(temp_dir / "missing.bin")
content = AssetContent(path=missing_path, hash=None, size_bytes=7, mtime_ns=1)
session.add(content)
session.flush()
record = Asset(content_id=content.id, name="missing.bin")
session.add(record)
session.commit()
sync_prefixes_with_filesystem(session, prefixes=[str(temp_dir)])
session.commit()
session.expire_all()
content = session.get(AssetContent, content.id)
missing_tag = session.get(AssetTag, {"asset_id": record.id, "tag_name": "missing"})
assert content is not None and content.is_missing is True
assert missing_tag is not None and missing_tag.origin == "automatic"
def test_enrichment_keeps_equal_hash_contents_distinct(session, temp_dir: Path):
path = temp_dir / "new.bin"
path.write_bytes(b"same bytes")
snapshot = snapshot_hash(str(path))
assert snapshot is not None
digest, _ = snapshot
content = AssetContent(
path=str(path),
hash=None,
size_bytes=path.stat().st_size,
mtime_ns=path.stat().st_mtime_ns,
)
existing = AssetContent(
path=str(temp_dir / "existing.bin"),
hash=to_stored_hash(digest),
size_bytes=path.stat().st_size,
mtime_ns=path.stat().st_mtime_ns,
)
session.add_all((content, existing))
session.flush()
record = Asset(content_id=content.id, name=path.name)
existing_record = Asset(content_id=existing.id, name="existing.bin")
session.add_all((record, existing_record))
session.commit()
enriched = enrich_asset(
session,
file_path=str(path),
content_id=content.id,
record_id=record.id,
extract_metadata=False,
compute_hash=True,
)
assert enriched is True
assert session.get(AssetContent, content.id).hash == to_stored_hash(digest)
assert session.get(Asset, record.id).content_id == content.id
assert session.get(Asset, existing_record.id).content_id == existing.id

View File

@@ -0,0 +1,214 @@
import os
from collections.abc import Callable
from pathlib import Path
from unittest.mock import patch
import pytest
from sqlalchemy import select
from sqlalchemy.exc import IntegrityError
from sqlalchemy.orm import Session
from app.assets.database.models import Asset, AssetContent
from app.assets.database.queries import create_content, create_record, delete_record
from app.assets.scanner import SeedAssetSpec, seed_asset_specs
from app.assets.services.snapshot_hash import snapshot_hash
def _spec(path: Path) -> SeedAssetSpec:
stat_result = path.stat()
return {
"abs_path": str(path),
"size_bytes": stat_result.st_size,
"mtime_ns": stat_result.st_mtime_ns,
"info_name": path.name,
"tags": ["input"],
"fname": path.name,
"metadata": None,
"mime_type": None,
"job_id": None,
}
def _specs_with_vanished_path(root: Path) -> tuple[list[SeedAssetSpec], Path]:
paths = [root / name for name in ("first.bin", "vanished.bin", "last.bin")]
for path in paths:
_ = path.write_bytes(path.name.encode())
return [_spec(path) for path in paths], paths[1]
def _record_count(session: Session) -> int:
return len(session.scalars(select(Asset)).all())
def test_seed_persists_remaining_specs_when_path_vanishes_before_restat(
session: Session, temp_dir: Path
) -> None:
specs, vanished_path = _specs_with_vanished_path(temp_dir)
vanished_path.unlink()
created = seed_asset_specs(session, specs)
session.commit()
assert created == 2
assert _record_count(session) == 2
def test_seed_persists_remaining_specs_when_path_vanishes_during_recovery_hash(
session: Session, temp_dir: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
specs, vanished_path = _specs_with_vanished_path(temp_dir)
def _hash_or_raise(path: str) -> str | None:
if path == str(vanished_path):
vanished_path.unlink()
raise OSError("file vanished during recovery")
return snapshot_hash(path)
monkeypatch.setattr("app.assets.scanner_changes.snapshot_hash", _hash_or_raise)
with patch("app.assets.scanner.mode.hashing_enabled", return_value=True):
created = seed_asset_specs(session, specs)
session.commit()
assert created == 2
assert _record_count(session) == 2
def _delete_before_restat(_monkeypatch: pytest.MonkeyPatch, path: Path) -> None:
path.unlink()
def _delete_during_recovery(monkeypatch: pytest.MonkeyPatch, path: Path) -> None:
def _hash_or_raise(candidate_path: str) -> str | None:
if candidate_path == str(path):
path.unlink()
raise OSError("file vanished during recovery")
return snapshot_hash(candidate_path)
monkeypatch.setattr("app.assets.scanner_changes.snapshot_hash", _hash_or_raise)
@pytest.mark.parametrize(
"delete_path",
[_delete_before_restat, _delete_during_recovery],
ids=["before-restat", "during-recovery"],
)
def test_seed_logs_once_for_each_vanished_path(
session: Session,
temp_dir: Path,
monkeypatch: pytest.MonkeyPatch,
caplog: pytest.LogCaptureFixture,
delete_path: Callable[[pytest.MonkeyPatch, Path], None],
) -> None:
specs, vanished_path = _specs_with_vanished_path(temp_dir)
delete_path(monkeypatch, vanished_path)
with patch("app.assets.scanner.mode.hashing_enabled", return_value=True):
_ = seed_asset_specs(session, specs)
session.commit()
messages = [
record.getMessage()
for record in caplog.records
if str(vanished_path) in record.getMessage()
]
assert messages == [f"Skipping vanished asset during scan: {vanished_path}"]
def test_seed_isolates_a_poisoned_spec_and_persists_the_specs_around_it(
session: Session, temp_dir: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
specs, poisoned_path = _specs_with_vanished_path(temp_dir)
def _create_record_or_raise(
session_arg: Session,
*,
content_id: str,
name: str,
mime_type: str | None,
job_id: str | None,
loader_path: str | None,
tags: list[str],
) -> Asset:
if name == poisoned_path.name:
raise IntegrityError("forced record creation failure", {}, ValueError())
return create_record(
session_arg,
content_id,
name,
mime_type,
job_id,
loader_path,
tags,
)
monkeypatch.setattr("app.assets.scanner.create_record", _create_record_or_raise)
created = seed_asset_specs(session, specs)
session.commit()
assert created == 2
assert _record_count(session) == 2
assert {record.name for record in session.scalars(select(Asset))} == {
"first.bin",
"last.bin",
}, (
"the batch shares one transaction, so a bare rollback would erase the spec BEFORE "
f"the poisoned one; both neighbours of {poisoned_path.name} must survive"
)
def test_seed_persists_fresh_stat_after_spec_was_built(
session: Session, temp_dir: Path
) -> None:
path = temp_dir / "changed-after-walk.bin"
path.write_bytes(b"walk-time")
spec = _spec(path)
path.write_bytes(b"fresh-seed-time-content")
fresh_mtime_ns = spec["mtime_ns"] + 1_000_000
os.utime(path, ns=(fresh_mtime_ns, fresh_mtime_ns))
fresh_stat = path.stat()
created = seed_asset_specs(session, [spec])
session.commit()
persisted = session.scalar(
select(AssetContent).where(AssetContent.path == str(path))
)
assert created == 1
assert fresh_stat.st_size != spec["size_bytes"]
assert fresh_stat.st_mtime_ns != spec["mtime_ns"]
assert persisted is not None
assert persisted.size_bytes == fresh_stat.st_size
assert persisted.mtime_ns == fresh_stat.st_mtime_ns
def test_seed_record_failure_preserves_retained_live_content(
session: Session, temp_dir: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
path = temp_dir / "retained.bin"
path.write_bytes(b"retained-live-content")
spec = _spec(path)
content = create_content(
session,
path=str(path),
size_bytes=spec["size_bytes"],
mtime_ns=spec["mtime_ns"],
)
record = create_record(session, content.id, path.name)
session.commit()
retained_content_id = content.id
delete_record(session, record.id)
session.commit()
def _raise_record_creation(*_args, **_kwargs):
raise RuntimeError("forced record creation failure")
monkeypatch.setattr("app.assets.scanner.create_record", _raise_record_creation)
with pytest.raises(RuntimeError, match="forced record creation failure"):
seed_asset_specs(session, [spec])
session.rollback()
assert session.get(AssetContent, retained_content_id) is not None

View File

@@ -0,0 +1,356 @@
import os
from datetime import datetime
from unittest.mock import patch
import pytest
from aiohttp import web
from aiohttp.test_utils import make_mocked_request
from sqlalchemy import select, update
from app.assets.api import routes
from app.assets.api.routes import _build_asset_response
from app.assets.database.models import Asset
from app.assets.database.queries.records import (
RecordPageSpec,
create_content,
create_record,
get_record_by_id,
list_records_page,
mark_content_missing,
)
from app.assets.helpers import to_stored_hash
from app.assets.services.asset_management import (
asset_exists,
delete_asset_reference,
get_asset_detail,
resolve_asset_for_download,
resolve_hash_to_path,
)
from app.assets.services.lookup import lookup_for_view
from app.assets.services.schemas import AssetData, AssetDetailResult, ReferenceData
_TS = datetime(2024, 1, 1, 0, 0, 0)
def test_missing_content_record_metadata_200_with_tag_content_404(
mock_create_session, session, temp_dir
):
content = create_content(session, path=str(temp_dir / "gone.bin"))
record = create_record(session, content_id=content.id, name="gone.bin")
mark_content_missing(session, content.id)
session.commit()
record_id = record.id
detail = get_asset_detail(record_id)
assert detail is not None
assert "missing" in detail.tags
with pytest.raises(FileNotFoundError):
resolve_asset_for_download(record_id)
def test_missing_record_is_listed_but_never_served(
mock_create_session, session, temp_dir
):
path = temp_dir / "listed_gone.bin"
path.write_bytes(b"x")
digest = to_stored_hash("a" * 64)
content = create_content(session, path=str(path), hash=digest)
record = create_record(session, content_id=content.id, name="listed_gone.bin")
mark_content_missing(session, content.id)
session.commit()
record_id = record.id
results, tag_map, total = list_records_page(session, RecordPageSpec())
assert any(r.id == record_id for r in results)
assert total == 1
assert "missing" in tag_map.get(record_id, [])
assert lookup_for_view(session, digest) is None
with pytest.raises(FileNotFoundError):
resolve_asset_for_download(record_id)
def test_deleted_record_metadata_and_content_404(mock_create_session, session):
record_id = "00000000-0000-0000-0000-000000000000"
assert get_asset_detail(record_id) is None
with pytest.raises(ValueError, match="not found"):
resolve_asset_for_download(record_id)
def test_content_404_when_file_absent_not_marked_missing(
mock_create_session, session, temp_dir
):
path = temp_dir / "absent.bin"
path.write_bytes(b"x")
content = create_content(session, path=str(path))
record = create_record(session, content_id=content.id, name="absent.bin")
session.commit()
record_id = record.id
path.unlink()
detail = get_asset_detail(record_id)
assert detail is not None
assert "missing" not in detail.tags
with pytest.raises(FileNotFoundError):
resolve_asset_for_download(record_id)
def test_content_fail_closed_no_sibling_fallback(mock_create_session, session, temp_dir):
digest = "d" * 64
missing_path = temp_dir / "missing.bin"
live_path = temp_dir / "live.bin"
live_path.write_bytes(b"live")
missing_content = create_content(session, path=str(missing_path), hash=digest)
live_content = create_content(session, path=str(live_path), hash=digest)
record = create_record(session, content_id=missing_content.id, name="missing.bin")
mark_content_missing(session, missing_content.id)
session.commit()
record_id = record.id
assert live_content.is_missing is False
assert os.path.isfile(live_path)
with pytest.raises(FileNotFoundError):
resolve_asset_for_download(record_id)
def test_resolve_hash_to_path_refuses_temp_content(mock_create_session, session, temp_dir):
digest = "e" * 64
f = temp_dir / "temp_only.bin"
f.write_bytes(b"temp")
create_content(session, path=str(f), hash=to_stored_hash(digest))
session.commit()
with patch("app.assets.services.lookup.is_temp_path", return_value=True):
result = resolve_hash_to_path(f"blake3:{digest}")
assert result is None
def test_resolve_hash_to_path_uses_newest_record_for_name_and_content_type(
mock_create_session, session, temp_dir
):
digest = "c" * 64
f = temp_dir / "shared.bin"
f.write_bytes(b"data")
content = create_content(session, path=str(f), hash=to_stored_hash(digest), size_bytes=f.stat().st_size)
older = create_record(
session,
content_id=content.id,
name="older.txt",
mime_type="text/plain",
)
newer = create_record(
session,
content_id=content.id,
name="newer.png",
mime_type="image/png",
)
older.created_at = _TS
older.updated_at = _TS
older.last_access_time = _TS
newer.created_at = _TS.replace(microsecond=1)
newer.updated_at = _TS.replace(microsecond=1)
newer.last_access_time = _TS.replace(microsecond=1)
session.commit()
result = resolve_hash_to_path(f"blake3:{digest}")
assert result is not None
assert result.download_name == "newer.png"
assert result.content_type == "image/png"
def test_resolve_hash_to_path_serves_content_with_zero_records(
mock_create_session, session, temp_dir
):
digest = "d" * 64
f = temp_dir / "orphan.png"
f.write_bytes(b"data")
content = create_content(session, path=str(f), hash=to_stored_hash(digest), size_bytes=f.stat().st_size)
record = create_record(
session,
content_id=content.id,
name="renamed.txt",
mime_type="text/plain",
)
session.commit()
assert delete_asset_reference(record.id) is True
session.expire_all()
assert lookup_for_view(session, f"blake3:{digest}") is not None
assert session.scalars(
select(Asset).where(Asset.content_id == content.id)
).all() == []
result = resolve_hash_to_path(f"blake3:{digest}")
assert result is not None
assert result.download_name == "orphan.png"
assert result.content_type == "image/png"
def test_resolve_hash_to_path_unknown_hash(mock_create_session):
assert resolve_hash_to_path("blake3:" + "f" * 64) is None
def test_content_read_updates_last_access_time(
mock_create_session, session, temp_dir
):
f = temp_dir / "read.bin"
f.write_bytes(b"data")
content = create_content(session, path=str(f))
record = create_record(session, content_id=content.id, name="read.bin")
session.commit()
record_id = record.id
session.execute(
update(Asset).where(Asset.id == record_id).values(last_access_time=None)
)
session.commit()
before_record = get_record_by_id(session, record_id)
assert before_record is not None
before = before_record.last_access_time
assert before is None
resolve_asset_for_download(record_id)
session.expire_all()
after_record = get_record_by_id(session, record_id)
assert after_record is not None
after = after_record.last_access_time
assert after is not None
def test_view_hash_read_updates_last_access_time(
mock_create_session, session, temp_dir
):
digest = "a" * 64
f = temp_dir / "view.bin"
f.write_bytes(b"view")
content = create_content(session, path=str(f), hash=to_stored_hash(digest), size_bytes=f.stat().st_size)
record = create_record(session, content_id=content.id, name="view.bin")
session.commit()
record_id = record.id
session.execute(
update(Asset).where(Asset.id == record_id).values(last_access_time=None)
)
session.commit()
result = resolve_hash_to_path(f"blake3:{digest}")
assert result is not None
session.expire_all()
record_after = get_record_by_id(session, record_id)
assert record_after is not None
assert record_after.last_access_time is not None
def test_lookup_for_view_returns_none_for_temp_content(session, temp_dir):
digest = "b" * 64
f = temp_dir / "view_temp.bin"
f.write_bytes(b"temp")
create_content(session, path=str(f), hash=to_stored_hash(digest))
session.commit()
with patch("app.assets.services.lookup.is_temp_path", return_value=True):
assert lookup_for_view(session, to_stored_hash(digest)) is None
def test_asset_exists_false_for_temp_content(mock_create_session, session, temp_dir):
digest = "c" * 64
f = temp_dir / "exists_temp.bin"
f.write_bytes(b"temp")
create_content(session, path=str(f), hash=to_stored_hash(digest))
session.commit()
with patch("app.assets.services.lookup.is_temp_path", return_value=True):
assert asset_exists(f"blake3:{digest}") is False
@pytest.mark.asyncio
async def test_head_hash_route_404_for_temp_content(
mock_create_session, session, temp_dir, monkeypatch
):
digest = "d" * 64
f = temp_dir / "head_temp.bin"
f.write_bytes(b"temp")
create_content(session, path=str(f), hash=to_stored_hash(digest))
session.commit()
monkeypatch.setattr(routes, "_ASSETS_ENABLED", True)
with patch("app.assets.services.lookup.is_temp_path", return_value=True):
response = await routes.head_asset_by_hash(
make_mocked_request(
"HEAD",
f"/api/assets/hash/blake3:{digest}",
match_info={"hash": f"blake3:{digest}"},
)
)
assert isinstance(response, web.Response)
assert response.status == 404
def test_resolve_hash_to_path_temp_does_not_bump_last_access_time(
mock_create_session, session, temp_dir
):
digest = "f" * 64
f = temp_dir / "noaccess_temp.bin"
f.write_bytes(b"temp")
content = create_content(session, path=str(f), hash=to_stored_hash(digest))
record = create_record(session, content_id=content.id, name="noaccess_temp.bin")
session.commit()
record_id = record.id
session.execute(
update(Asset).where(Asset.id == record_id).values(last_access_time=None)
)
session.commit()
with patch("app.assets.services.lookup.is_temp_path", return_value=True):
result = resolve_hash_to_path(f"blake3:{digest}")
assert result is None
session.expire_all()
record_after = get_record_by_id(session, record_id)
assert record_after is not None
assert record_after.last_access_time is None
@pytest.fixture
def sandboxed_comfy_roots(tmp_path):
with patch("app.assets.services.path_utils.folder_paths") as fp:
fp.get_input_directory.return_value = str(tmp_path / "input")
fp.get_output_directory.return_value = str(tmp_path / "output")
fp.get_temp_directory.return_value = str(tmp_path / "temp")
fp.models_dir = str(tmp_path / "models")
yield tmp_path
def test_temp_asset_preview_url_still_resolves_type_temp(sandboxed_comfy_roots):
name = "ComfyUI_temp_abcde_00001_.png"
result = AssetDetailResult(
ref=ReferenceData(
id="ref-temp",
name=name,
file_path=str(sandboxed_comfy_roots / "temp" / name),
loader_path=None,
user_metadata=None,
preview_id=None,
created_at=_TS,
updated_at=_TS,
last_access_time=_TS,
),
asset=AssetData(hash="blake3:abc", size_bytes=1024, mime_type="image/png"),
tags=[],
)
resp = _build_asset_response(result, {})
assert resp.preview_url == f"/api/view?type=temp&filename={name}"

View File

@@ -0,0 +1,90 @@
import builtins
import os
from collections.abc import Callable
from pathlib import Path
import pytest
from blake3 import blake3
from app.assets.services.snapshot_hash import snapshot_hash
class _MutatingReader:
def __init__(self, file, mutate: Callable[[], None], mutate_after_close: bool) -> None:
self._file = file
self._mutate = mutate
self._mutate_after_close = mutate_after_close
self._did_mutate = False
def __enter__(self):
self._file.__enter__()
return self
def __exit__(self, *args):
result = self._file.__exit__(*args)
if self._mutate_after_close and not self._did_mutate:
self._did_mutate = True
self._mutate()
return result
def fileno(self) -> int:
return self._file.fileno()
def read(self, size: int = -1) -> bytes:
result = self._file.read(size)
if not self._mutate_after_close and not self._did_mutate:
self._did_mutate = True
self._mutate()
return result
def test_snapshot_hash_returns_digest_for_quiescent_file(tmp_path: Path) -> None:
payload = b"quiescent" * 1024
path = tmp_path / "asset.bin"
path.write_bytes(payload)
result = snapshot_hash(str(path), chunk_size=64)
assert result is not None
digest, stat_result = result
assert digest == blake3(payload).hexdigest()
assert isinstance(stat_result, os.stat_result)
assert stat_result.st_size == len(payload)
@pytest.mark.parametrize("mutation", ["bytes", "replace", "unlink", "truncate", "append"])
def test_snapshot_hash_returns_none_when_file_drifts(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch, mutation: str
) -> None:
path = tmp_path / "asset.bin"
path.write_bytes(b"original" * 1024)
original_open = builtins.open
def mutate() -> None:
match mutation:
case "bytes":
path.write_bytes(b"changed" * 1024)
case "replace":
replacement = tmp_path / "replacement.bin"
replacement.write_bytes(b"replacement")
os.replace(replacement, path)
case "unlink":
path.unlink()
case "truncate":
path.write_bytes(b"")
case "append":
with original_open(path, "ab") as output:
output.write(b"more")
case unreachable:
raise AssertionError(f"unexpected mutation {unreachable}")
def open_with_mutation(*args, **kwargs):
return _MutatingReader(
original_open(*args, **kwargs),
mutate,
mutation in {"replace", "unlink"},
)
monkeypatch.setattr(builtins, "open", open_with_mutation)
assert snapshot_hash(str(path), chunk_size=64) is None

View File

@@ -0,0 +1,393 @@
from __future__ import annotations
import os
from collections.abc import Iterator
from contextlib import contextmanager
from dataclasses import dataclass
from pathlib import Path
from unittest.mock import patch
import sqlalchemy as sa
from sqlalchemy import select
from sqlalchemy.orm import Session
import pytest
from app.assets.database.models import Asset, AssetContent
from app.assets.database.queries import create_content, create_record
from app.assets.database.queries.records import (
RecordPageSpec,
fetch_record_tags,
list_records_page,
)
from app.assets.helpers import to_stored_hash
from app.assets.scanner import get_unenriched_assets_for_roots
from app.assets.scanner_changes import (
clear_pending_verifications,
detect_content_change,
drain_pending_verifications,
)
from app.assets.services.hash_mode_state import (
clear_transition_queue,
drain_transition_queue,
enqueue_transition_work,
)
from app.assets.services.lookup import lookup_for_view
from app.assets.services.snapshot_hash import snapshot_hash
@dataclass(frozen=True, slots=True)
class _FakeStat:
st_size: int
st_mtime_ns: int
@contextmanager
def _reuse_session(session: Session) -> Iterator[Session]:
yield session
def _raw_system_metadata(session: Session, record_id: str) -> object:
return session.execute(
sa.text("SELECT system_metadata FROM assets WHERE id = :id"),
{"id": record_id},
).scalar()
def _candidates_under(session: Session, temp_dir: Path, *, compute_hashes: bool) -> set[str]:
with (
patch("app.assets.scanner.create_session", lambda: _reuse_session(session)),
patch(
"app.assets.scanner.get_scan_prefixes_for_root",
return_value=[str(temp_dir)],
),
):
rows = get_unenriched_assets_for_roots(("models",), compute_hashes=compute_hashes)
return {row.record_id for row in rows}
@pytest.fixture(autouse=True)
def _transition_queue_isolation() -> Iterator[None]:
clear_transition_queue()
yield
clear_transition_queue()
@pytest.fixture(autouse=True)
def _input_base_is_temp_dir(temp_dir: Path, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr("folder_paths.get_input_directory", lambda: str(temp_dir))
@pytest.fixture(autouse=True)
def _pending_verification_isolation() -> Iterator[None]:
clear_pending_verifications()
yield
clear_pending_verifications()
def _seed_hashed_content(session: Session, path: Path, data: bytes) -> AssetContent:
path.write_bytes(data)
snapshot = snapshot_hash(str(path))
assert snapshot is not None
digest, verified_stat = snapshot
return create_content(
session,
str(path),
to_stored_hash(digest),
verified_stat.st_size,
verified_stat.st_mtime_ns,
)
def _bump_mtime(path: Path) -> os.stat_result:
before = path.stat()
bumped = before.st_mtime_ns + 5_000_000_000
os.utime(path, ns=(bumped, bumped))
after = path.stat()
assert after.st_mtime_ns != before.st_mtime_ns
assert after.st_size == before.st_size
return after
def _rewrite_same_size(path: Path, data: bytes) -> os.stat_result:
before = path.stat()
assert len(data) == before.st_size
assert data != path.read_bytes()
path.write_bytes(data)
bumped = before.st_mtime_ns + 5_000_000_000
os.utime(path, ns=(bumped, bumped))
after = path.stat()
assert after.st_size == before.st_size
assert after.st_mtime_ns != before.st_mtime_ns
return after
def test_split_record_is_enrich_candidate_in_hash_mode(
session: Session, temp_dir: Path
) -> None:
path = temp_dir / "split.safetensors"
content = create_content(session, str(path), hash="blake3:deadbeef")
record = create_record(session, content.id, path.name)
session.commit()
record_id = record.id
candidates = _candidates_under(session, temp_dir, compute_hashes=True)
assert record_id in candidates
def test_enriched_record_is_not_enrich_candidate_in_hash_mode(
session: Session, temp_dir: Path
) -> None:
path = temp_dir / "enriched.safetensors"
content = create_content(session, str(path), hash="blake3:deadbeef")
record = create_record(
session, content.id, path.name, system_metadata={"architecture": "flux"}
)
session.commit()
record_id = record.id
candidates = _candidates_under(session, temp_dir, compute_hashes=True)
assert record_id not in candidates
def test_same_size_mtime_bump_does_not_split(session: Session, temp_dir: Path) -> None:
path = temp_dir / "touched.safetensors"
content = create_content(session, str(path), hash=None, size_bytes=100, mtime_ns=1000)
record = create_record(
session, content.id, path.name, tags=["keepme"], system_metadata={"k": "v"}
)
session.commit()
content_id, record_id = content.id, record.id
detect_content_change(
session, content, _FakeStat(st_size=100, st_mtime_ns=2000), hashing_is_enabled=False
)
session.commit()
session.expire_all()
live = session.get(AssetContent, content_id)
assert live is not None and live.is_missing is False
rows_at_path = list(
session.scalars(select(AssetContent).where(AssetContent.path == str(path)))
)
assert len(rows_at_path) == 1
surviving = session.get(Asset, record_id)
assert surviving.system_metadata == {"k": "v"}
tags = fetch_record_tags(session, record_id)
assert "keepme" in tags
assert "missing" not in tags
assert live.mtime_ns == 2000
assert live.size_bytes == 100
def test_accepted_mtime_bump_drops_the_unverifiable_hash(
session: Session, temp_dir: Path
) -> None:
path = temp_dir / "synced.safetensors"
content = _seed_hashed_content(session, path, b"rsynced bytes")
record = create_record(
session, content.id, path.name, tags=["keepme"], system_metadata={"k": "v"}
)
session.commit()
content_id, record_id, stored_hash = content.id, record.id, content.hash
assert lookup_for_view(session, stored_hash) is not None
observed = _bump_mtime(path)
detect_content_change(session, content, observed, hashing_is_enabled=False)
session.commit()
session.expire_all()
live = session.get(AssetContent, content_id)
assert live.hash is None
assert lookup_for_view(session, stored_hash) is None
assert live.is_missing is False
assert live.mtime_ns == observed.st_mtime_ns
assert live.size_bytes == observed.st_size
surviving = session.get(Asset, record_id)
assert surviving is not None and surviving.content_id == content_id
assert surviving.system_metadata == {"k": "v"}
assert "keepme" in fetch_record_tags(session, record_id)
listed, _, _ = list_records_page(session, RecordPageSpec(limit=100))
assert record_id in {row.id for row in listed}
def test_same_size_content_change_is_never_served_under_the_old_hash(
session: Session, temp_dir: Path
) -> None:
path = temp_dir / "overwritten.safetensors"
content = _seed_hashed_content(session, path, b"AAAA")
create_record(session, content.id, path.name, tags=["keepme"])
session.commit()
old_hash = content.hash
assert lookup_for_view(session, old_hash) is not None
observed = _rewrite_same_size(path, b"BBBB")
detect_content_change(session, content, observed, hashing_is_enabled=False)
session.commit()
session.expire_all()
assert path.read_bytes() == b"BBBB"
assert lookup_for_view(session, old_hash) is None
def test_accepted_mtime_bump_is_not_re_detected_by_the_next_scan(
session: Session, temp_dir: Path
) -> None:
path = temp_dir / "resynced.safetensors"
content = _seed_hashed_content(session, path, b"cloud synced bytes")
create_record(session, content.id, path.name, tags=["keepme"])
session.commit()
content_id = content.id
detect_content_change(session, content, _bump_mtime(path), hashing_is_enabled=False)
session.commit()
detect_content_change(session, content, path.stat(), hashing_is_enabled=True)
assert drain_pending_verifications(session) == 0
detect_content_change(session, content, path.stat(), hashing_is_enabled=False)
session.commit()
session.expire_all()
rows_at_path = list(
session.scalars(select(AssetContent).where(AssetContent.path == str(path)))
)
assert len(rows_at_path) == 1
assert rows_at_path[0].id == content_id
assert rows_at_path[0].is_missing is False
def test_dropped_hash_is_refilled_in_place_by_a_later_hash_mode_pass(
session: Session, temp_dir: Path
) -> None:
path = temp_dir / "refilled.safetensors"
content = _seed_hashed_content(session, path, b"cloud synced bytes")
record = create_record(
session, content.id, path.name, tags=["keepme"], system_metadata={"k": "v"}
)
session.commit()
content_id, record_id = content.id, record.id
assert record_id not in _candidates_under(session, temp_dir, compute_hashes=True)
detect_content_change(session, content, _bump_mtime(path), hashing_is_enabled=False)
session.commit()
session.expire_all()
assert session.get(AssetContent, content_id).hash is None
assert record_id in _candidates_under(session, temp_dir, compute_hashes=True)
enqueue_transition_work(session, "off_to_on")
drain_transition_queue(session)
session.commit()
session.expire_all()
snapshot = snapshot_hash(str(path))
assert snapshot is not None
live = session.get(AssetContent, content_id)
assert live.is_missing is False
assert live.hash == to_stored_hash(snapshot[0])
assert lookup_for_view(session, live.hash).id == content_id
rows_at_path = list(
session.scalars(select(AssetContent).where(AssetContent.path == str(path)))
)
assert len(rows_at_path) == 1
assert "keepme" in fetch_record_tags(session, record_id)
assert session.get(Asset, record_id).system_metadata == {"k": "v"}
assert record_id not in _candidates_under(session, temp_dir, compute_hashes=True)
def test_mtime_and_size_change_splits_with_null_metadata(
session: Session, temp_dir: Path
) -> None:
path = temp_dir / "grown.safetensors"
content = create_content(session, str(path), hash=None, size_bytes=100, mtime_ns=1000)
create_record(
session, content.id, path.name, tags=["oldtag"], system_metadata={"k": "v"}
)
session.commit()
old_content_id = content.id
detect_content_change(
session, content, _FakeStat(st_size=200, st_mtime_ns=2000), hashing_is_enabled=False
)
session.commit()
session.expire_all()
assert session.get(AssetContent, old_content_id).is_missing is True
live = session.scalar(
select(AssetContent).where(
AssetContent.path == str(path), AssetContent.is_missing.is_(False)
)
)
assert live is not None and live.id != old_content_id
assert live.size_bytes == 200 and live.mtime_ns == 2000
new_record = session.scalar(select(Asset).where(Asset.content_id == live.id))
assert new_record is not None
assert new_record.system_metadata is None
assert _raw_system_metadata(session, new_record.id) is None
assert "oldtag" not in fetch_record_tags(session, new_record.id)
assert new_record.id in _candidates_under(session, temp_dir, compute_hashes=True)
def test_mtime_unchanged_size_changed_does_not_split(
session: Session, temp_dir: Path
) -> None:
path = temp_dir / "weird.safetensors"
content = create_content(session, str(path), hash=None, size_bytes=100, mtime_ns=1000)
create_record(session, content.id, path.name, tags=["keepme"])
session.commit()
content_id = content.id
detect_content_change(
session, content, _FakeStat(st_size=999, st_mtime_ns=1000), hashing_is_enabled=False
)
session.commit()
session.expire_all()
assert session.get(AssetContent, content_id).is_missing is False
rows_at_path = list(
session.scalars(select(AssetContent).where(AssetContent.path == str(path)))
)
assert len(rows_at_path) == 1
def test_transition_drain_split_replacement_has_null_metadata(
session: Session, temp_dir: Path
) -> None:
path = temp_dir / "changed.bin"
path.write_bytes(b"old bytes")
old_snapshot = snapshot_hash(str(path))
assert old_snapshot is not None
old_digest, _ = old_snapshot
stat = path.stat()
old_content = create_content(
session, str(path), to_stored_hash(old_digest), stat.st_size, stat.st_mtime_ns
)
old_content_id = old_content.id
create_record(
session, old_content_id, "changed.bin", tags=["oldtag"], system_metadata={"k": "v"}
)
path.write_bytes(b"different new bytes")
enqueue_transition_work(session, "off_to_on")
drain_transition_queue(session)
session.commit()
session.expire_all()
assert session.get(AssetContent, old_content_id).is_missing is True
live = session.scalar(
select(AssetContent).where(
AssetContent.path == str(path), AssetContent.is_missing.is_(False)
)
)
assert live is not None and live.id != old_content_id
new_record = session.scalar(select(Asset).where(Asset.content_id == live.id))
assert new_record is not None
assert new_record.system_metadata is None
assert _raw_system_metadata(session, new_record.id) is None
assert "oldtag" not in fetch_record_tags(session, new_record.id)

View File

@@ -0,0 +1,282 @@
from __future__ import annotations
import json
import os
import uuid
from contextlib import contextmanager
from unittest.mock import AsyncMock
import pytest
from aiohttp import web
from aiohttp.test_utils import make_mocked_request
from sqlalchemy import func, select
from sqlalchemy.orm import Session as SASession
from app.assets import mode
from app.assets.api import routes
import app.assets.mode as mode_module
import folder_paths
from app.assets.database.models import Asset, AssetContent
from app.assets.database.queries.records import create_record
from app.assets.scanner import enrich_asset
from app.assets.scanner_changes import recover_missing_content
from app.assets.services import asset_management, ingest
from app.assets.services.asset_management import get_asset_detail
from app.assets.services.ingest import (
HashMismatchError,
upload_from_temp_path,
)
from app.assets.services.snapshot_hash import snapshot_hash
@pytest.fixture
def hashing_on():
class FakeArgs:
enable_asset_hashing = True
class DisabledArgs:
enable_asset_hashing = False
mode_module.init(FakeArgs())
yield
mode_module.init(DisabledArgs())
def _write_temp(content: bytes) -> str:
uploads_root = os.path.join(
folder_paths.get_temp_directory(), "uploads", uuid.uuid4().hex
)
os.makedirs(uploads_root, exist_ok=True)
path = os.path.join(uploads_root, ".upload.part")
with open(path, "wb") as file:
file.write(content)
return path
def test_upload_stores_prefixed_hash_expected_hash_succeeds_and_reuses_content(
mock_create_session, hashing_on
):
content_bytes = b"prefixed-stored-upload-bytes"
temp1 = _write_temp(content_bytes)
snapshot = snapshot_hash(temp1)
assert snapshot is not None
digest, _ = snapshot
expected = f"blake3:{digest}"
assert expected.startswith("blake3:")
temp2 = _write_temp(content_bytes)
try:
r1 = upload_from_temp_path(
temp_path=temp1,
name="pref.bin",
tags=["output"],
client_filename="pref.bin",
expected_hash=expected,
)
assert r1.created_new is True
with mock_create_session() as session:
live = list(
session.scalars(
select(AssetContent).where(AssetContent.is_missing.is_(False))
)
)
assert len(live) == 1
assert live[0].hash == expected
assert live[0].hash.startswith("blake3:")
assert r1.asset.hash == expected
r2 = upload_from_temp_path(
temp_path=temp2,
name="pref.bin",
tags=["output"],
client_filename="pref.bin",
expected_hash=expected,
)
assert r2.created_new is True
assert r2.ref.id != r1.ref.id, "every upload is its own delivery record"
with mock_create_session() as session:
assert session.scalar(select(func.count()).select_from(AssetContent)) == 1
assert session.scalar(select(func.count()).select_from(Asset)) == 2
finally:
for path in (temp1, temp2):
if os.path.exists(path):
os.unlink(path)
def test_upload_expected_hash_mismatch_still_rejected(mock_create_session, hashing_on):
temp = _write_temp(b"mismatch-bytes")
try:
with pytest.raises(HashMismatchError):
upload_from_temp_path(
temp_path=temp,
name="mismatch.bin",
tags=["output"],
client_filename="mismatch.bin",
expected_hash=f"blake3:{'0' * 64}",
)
with mock_create_session() as session:
assert session.scalar(select(func.count()).select_from(Asset)) == 0
assert session.scalar(select(func.count()).select_from(AssetContent)) == 0
finally:
if os.path.exists(temp):
os.unlink(temp)
def test_enrich_fills_deferred_hash_prefixed(session, temp_dir):
path = temp_dir / "deferred.bin"
path.write_bytes(b"deferred output bytes")
stat = path.stat()
content = AssetContent(
path=str(path),
hash=None,
size_bytes=stat.st_size,
mtime_ns=stat.st_mtime_ns,
)
session.add(content)
session.flush()
record = create_record(session, content.id, path.name)
session.commit()
enriched = enrich_asset(
session,
file_path=str(path),
content_id=content.id,
record_id=record.id,
extract_metadata=False,
compute_hash=True,
)
assert enriched is True
stored = session.get(AssetContent, content.id).hash
assert stored is not None
assert stored.startswith("blake3:")
snapshot = snapshot_hash(str(path))
assert snapshot is not None
digest, _ = snapshot
assert stored == f"blake3:{digest}"
def test_recovery_matches_prefixed_stored_hash(session, temp_dir):
path = temp_dir / "recover.bin"
original_bytes = b"recover-me-bytes"
path.write_bytes(original_bytes)
snapshot = snapshot_hash(str(path))
assert snapshot is not None
digest, _ = snapshot
stored = f"blake3:{digest}"
assert stored.startswith("blake3:")
content = AssetContent(
path=str(path),
hash=stored,
is_missing=True,
size_bytes=0,
mtime_ns=None,
)
session.add(content)
session.flush()
create_record(session, content.id, path.name)
session.commit()
path.unlink()
path.write_bytes(original_bytes)
stat = os.stat(str(path))
result = recover_missing_content(
session, str(path), stat, hashing_is_enabled=True
)
assert result == "recovered"
recovered = session.get(AssetContent, content.id)
assert recovered is not None
assert recovered.is_missing is False
assert recovered.hash == stored
@pytest.mark.asyncio
async def test_all_read_surfaces_agree_on_prefixed_hash(
db_engine, monkeypatch, hashing_on
):
@contextmanager
def _factory():
with SASession(db_engine) as sess:
yield sess
monkeypatch.setattr(routes, "create_session", lambda: SASession(db_engine))
monkeypatch.setattr(routes, "_ASSETS_ENABLED", True)
monkeypatch.setattr(mode, "hashing_enabled", lambda: True)
monkeypatch.setattr(ingest, "create_session", _factory)
monkeypatch.setattr(asset_management, "create_session", _factory)
content_bytes = b"one-asset-all-surfaces-agree"
temp = _write_temp(content_bytes)
snapshot = snapshot_hash(temp)
assert snapshot is not None
digest, _ = snapshot
expected = f"blake3:{digest}"
try:
upload_result = ingest.upload_from_temp_path(
temp_path=temp,
name="surf.bin",
tags=["output"],
client_filename="surf.bin",
)
asset_id = upload_result.ref.id
assert upload_result.asset is not None
assert upload_result.asset.hash == expected
upload_resp = routes._build_asset_response(upload_result, {})
assert upload_resp.hash == expected
detail = get_asset_detail(asset_id)
assert detail is not None
assert detail.asset is not None
assert detail.asset.hash == expected
get_resp = await routes.get_asset_route(
make_mocked_request(
"GET", f"/api/assets/{asset_id}", match_info={"id": asset_id}
)
)
assert isinstance(get_resp, web.Response)
assert isinstance(get_resp.body, bytes | bytearray)
get_body = json.loads(get_resp.body)
assert get_body["hash"] == expected
list_resp = await routes.list_assets_route(
make_mocked_request("GET", "/api/assets")
)
assert isinstance(list_resp, web.Response)
assert isinstance(list_resp.body, bytes | bytearray)
list_body = json.loads(list_resp.body)
item = next(a for a in list_body["assets"] if a["id"] == asset_id)
assert item["hash"] == expected
head_resp = await routes.head_asset_by_hash(
make_mocked_request(
"HEAD",
f"/api/assets/hash/{expected}",
match_info={"hash": expected},
)
)
assert isinstance(head_resp, web.Response)
assert head_resp.status == 200
fh_req = AsyncMock(spec=web.Request)
fh_req.json.return_value = {
"hash": expected,
"name": "fromhash.bin",
"tags": ["output"],
}
fh_resp = await routes.create_asset_from_hash_route(fh_req)
assert isinstance(fh_resp, web.Response)
assert fh_resp.status == 201
assert isinstance(fh_resp.body, bytes | bytearray)
fh_body = json.loads(fh_resp.body)
assert fh_body["hash"] == expected
finally:
if os.path.exists(temp):
os.unlink(temp)

View File

@@ -1,123 +1,17 @@
"""Tests for list_tag_histogram service function."""
from sqlalchemy.orm import Session
from app.assets.database.models import Asset, AssetReference
from app.assets.database.queries import ensure_tags_exist, add_tags_to_reference
from app.assets.helpers import get_utc_now
from app.assets.services.tagging import list_tag_histogram
from app.assets.database.queries import create_content, create_record, list_records_page
from app.assets.database.queries.records import RecordPageSpec
def _make_asset(session: Session, hash_val: str = "blake3:test") -> Asset:
asset = Asset(hash=hash_val, size_bytes=1024)
session.add(asset)
session.flush()
return asset
def test_tag_selection_counts_records_not_content_rows(session: Session) -> None:
content = create_content(session, "/models/shared.safetensors")
create_record(session, content.id, "one", tags=["model"])
create_record(session, content.id, "two", tags=["model"])
def _make_reference(
session: Session,
asset: Asset,
name: str = "test",
owner_id: str = "",
) -> AssetReference:
now = get_utc_now()
ref = AssetReference(
owner_id=owner_id,
name=name,
asset_id=asset.id,
created_at=now,
updated_at=now,
last_access_time=now,
records, _, _ = list_records_page(
session,
RecordPageSpec(all_tags=("model",)),
)
session.add(ref)
session.flush()
return ref
class TestListTagHistogram:
def test_returns_counts_for_all_tags(self, mock_create_session, session: Session):
ensure_tags_exist(session, ["alpha", "beta"])
a1 = _make_asset(session, "blake3:aaa")
r1 = _make_reference(session, a1, name="r1")
add_tags_to_reference(session, reference_id=r1.id, tags=["alpha", "beta"])
a2 = _make_asset(session, "blake3:bbb")
r2 = _make_reference(session, a2, name="r2")
add_tags_to_reference(session, reference_id=r2.id, tags=["alpha"])
session.commit()
result = list_tag_histogram()
assert result["alpha"] == 2
assert result["beta"] == 1
def test_empty_when_no_assets(self, mock_create_session, session: Session):
ensure_tags_exist(session, ["unused"])
session.commit()
result = list_tag_histogram()
assert result == {}
def test_include_tags_filter(self, mock_create_session, session: Session):
ensure_tags_exist(session, ["models", "loras", "input"])
a1 = _make_asset(session, "blake3:aaa")
r1 = _make_reference(session, a1, name="r1")
add_tags_to_reference(session, reference_id=r1.id, tags=["models", "loras"])
a2 = _make_asset(session, "blake3:bbb")
r2 = _make_reference(session, a2, name="r2")
add_tags_to_reference(session, reference_id=r2.id, tags=["input"])
session.commit()
result = list_tag_histogram(include_tags=["models"])
# Only r1 has "models", so only its tags appear
assert "models" in result
assert "loras" in result
assert "input" not in result
def test_exclude_tags_filter(self, mock_create_session, session: Session):
ensure_tags_exist(session, ["models", "loras", "input"])
a1 = _make_asset(session, "blake3:aaa")
r1 = _make_reference(session, a1, name="r1")
add_tags_to_reference(session, reference_id=r1.id, tags=["models", "loras"])
a2 = _make_asset(session, "blake3:bbb")
r2 = _make_reference(session, a2, name="r2")
add_tags_to_reference(session, reference_id=r2.id, tags=["input"])
session.commit()
result = list_tag_histogram(exclude_tags=["models"])
# r1 excluded, only r2's tags remain
assert "input" in result
assert "loras" not in result
def test_name_contains_filter(self, mock_create_session, session: Session):
ensure_tags_exist(session, ["alpha", "beta"])
a1 = _make_asset(session, "blake3:aaa")
r1 = _make_reference(session, a1, name="my_model.safetensors")
add_tags_to_reference(session, reference_id=r1.id, tags=["alpha"])
a2 = _make_asset(session, "blake3:bbb")
r2 = _make_reference(session, a2, name="picture.png")
add_tags_to_reference(session, reference_id=r2.id, tags=["beta"])
session.commit()
result = list_tag_histogram(name_contains="model")
assert "alpha" in result
assert "beta" not in result
def test_limit_caps_results(self, mock_create_session, session: Session):
tags = [f"tag{i}" for i in range(10)]
ensure_tags_exist(session, tags)
a = _make_asset(session, "blake3:aaa")
r = _make_reference(session, a, name="r1")
add_tags_to_reference(session, reference_id=r.id, tags=tags)
session.commit()
result = list_tag_histogram(limit=3)
assert len(result) == 3
assert len(records) == 2

View File

@@ -0,0 +1,107 @@
import json
from types import SimpleNamespace
import pytest
from app.assets.api import routes
_RECORD_ID = "00000000-0000-0000-0000-000000000001"
_SYSTEM_TAG_ENVELOPE = {
"error": {
"code": "SYSTEM_TAG_FORBIDDEN",
"message": "Tag 'missing' is system-managed and cannot be modified via the API",
"details": {"tag": "missing"},
}
}
class _JsonRequest:
def __init__(self, tags: list[str]) -> None:
self.match_info = {"id": _RECORD_ID}
self._payload = {"tags": tags}
async def json(self) -> dict[str, list[str]]:
return self._payload
class _UserManager:
def get_request_user_id(self, _request: _JsonRequest) -> str:
return "test-user"
@pytest.mark.asyncio
async def test_add_missing_tag_returns_400(monkeypatch):
monkeypatch.setattr(routes, "USER_MANAGER", _UserManager())
monkeypatch.setattr(
routes,
"apply_tags",
lambda **_kwargs: SimpleNamespace(added=[], already_present=[], total_tags=[]),
)
response = await routes.add_asset_tags.__wrapped__(_JsonRequest(["missing"]))
assert response.status == 400
assert json.loads(response.body) == _SYSTEM_TAG_ENVELOPE
@pytest.mark.asyncio
async def test_remove_missing_tag_returns_400(monkeypatch):
monkeypatch.setattr(routes, "USER_MANAGER", _UserManager())
monkeypatch.setattr(
routes,
"remove_tags",
lambda **_kwargs: SimpleNamespace(removed=[], not_present=[], total_tags=[]),
)
response = await routes.delete_asset_tags.__wrapped__(_JsonRequest(["missing"]))
assert response.status == 400
assert json.loads(response.body) == _SYSTEM_TAG_ENVELOPE
@pytest.mark.asyncio
async def test_other_tags_unaffected(monkeypatch):
monkeypatch.setattr(routes, "USER_MANAGER", _UserManager())
calls: list[tuple[str, list[str]]] = []
def apply_tags(**kwargs):
calls.append(("add", kwargs["tags"]))
return SimpleNamespace(
added=kwargs["tags"], already_present=[], total_tags=kwargs["tags"]
)
def remove_tags(**kwargs):
calls.append(("remove", kwargs["tags"]))
return SimpleNamespace(
removed=kwargs["tags"], not_present=[], total_tags=[], protected=[]
)
monkeypatch.setattr(routes, "apply_tags", apply_tags)
monkeypatch.setattr(routes, "remove_tags", remove_tags)
add_response = await routes.add_asset_tags.__wrapped__(_JsonRequest(["manual"]))
remove_response = await routes.delete_asset_tags.__wrapped__(
_JsonRequest(["manual"])
)
assert add_response.status == 200
assert remove_response.status == 200
assert calls == [("add", ["manual"]), ("remove", ["manual"])]
@pytest.mark.asyncio
async def test_remove_tags_response_exposes_protected_bucket(monkeypatch):
monkeypatch.setattr(routes, "USER_MANAGER", _UserManager())
monkeypatch.setattr(
routes,
"remove_tags",
lambda **_kwargs: SimpleNamespace(
removed=[], not_present=[], total_tags=["auto"], protected=["auto"]
),
)
response = await routes.delete_asset_tags.__wrapped__(_JsonRequest(["auto"]))
assert response.status == 200
body = json.loads(response.body)
assert body["protected"] == ["auto"]

View File

@@ -0,0 +1,166 @@
from contextlib import contextmanager
from datetime import datetime
from unittest.mock import patch
import pytest
from sqlalchemy import create_engine, select, update
from sqlalchemy.exc import IntegrityError
from sqlalchemy.orm import Session as SASession
import app.assets.database.queries.records as records_module
import app.assets.services.asset_management as asset_management_module
import app.assets.services.tagging as tagging_module
from app.assets.database.models import Asset, AssetTag, Base, Tag
from app.assets.database.queries.records import (
create_content,
create_record,
mark_content_missing,
)
from app.assets.services.asset_management import update_asset_metadata
from app.assets.services.tagging import apply_tags
STALE = datetime(2020, 1, 1, 0, 0, 0)
RACED = "raced"
UNRELATED = "unrelated"
def _file_backed_engine(db_path):
engine = create_engine(
f"sqlite:///{db_path}",
connect_args={"check_same_thread": False, "timeout": 0.2},
)
Base.metadata.create_all(engine)
return engine
def _seed_record(session, path: str, tags=None) -> str:
content = create_content(session, path)
record = create_record(session, content.id, "fixture", tags=tags)
session.commit()
session.execute(update(Asset).where(Asset.id == record.id).values(updated_at=STALE))
session.commit()
return record.id
def _updated_at(session, record_id: str) -> datetime:
session.expire_all()
return session.scalar(select(Asset.updated_at).where(Asset.id == record_id))
def _tag_names(session, record_id: str) -> list[str]:
return sorted(
session.scalars(select(AssetTag.tag_name).where(AssetTag.asset_id == record_id)).all()
)
def test_apply_tags_loses_a_tag_race_without_raising_and_reports_it_honestly(tmp_path):
engine = _file_backed_engine(tmp_path / "tag_race.db")
with SASession(engine) as seed_session:
record_id = _seed_record(seed_session, "/tmp/tag-race-fixture", tags=[UNRELATED])
def connection_a_commits_the_raced_tag() -> None:
with SASession(engine) as connection_a:
connection_a.add(Tag(name=RACED))
connection_a.flush()
connection_a.add(
AssetTag(asset_id=record_id, tag_name=RACED, origin="manual")
)
connection_a.commit()
fired: list[bool] = []
@contextmanager
def racing_session_factory():
with SASession(engine) as session_b:
real_add = session_b.add
def add_racing_the_winner(instance, *args, **kwargs):
if isinstance(instance, Tag) and instance.name == RACED and not fired:
fired.append(True)
connection_a_commits_the_raced_tag()
return real_add(instance, *args, **kwargs)
session_b.add = add_racing_the_winner
yield session_b
with patch("app.assets.services.tagging.create_session", racing_session_factory):
result = apply_tags(record_id, [RACED])
assert fired, "the interleave never fired; the test proves nothing"
assert result.added == [], (
"session B lost the race; it inserted nothing and must not claim it did"
)
assert result.already_present == [RACED], (
"already_present reports the REQUESTED tags that were already there — never "
"unrelated pre-existing tags the caller did not ask about"
)
assert sorted(result.total_tags) == [RACED, UNRELATED]
with SASession(engine) as check:
assert _tag_names(check, record_id) == [RACED, UNRELATED]
assert _updated_at(check, record_id) == STALE, (
"a race loser changed no link, so it is not an edit and must not move updated_at"
)
def test_ensure_tag_link_reraises_when_the_parent_asset_is_missing(db_engine_fk):
with SASession(db_engine_fk) as session:
session.add(Tag(name="orphan-link"))
session.flush()
with pytest.raises(IntegrityError):
records_module.ensure_tag_link(
session,
asset_id="no-such-asset",
tag_name="orphan-link",
origin="manual",
)
def _drive_create_record(session, record_id):
content = create_content(session, "/tmp/four-sites-create-record")
create_record(session, content.id, "created", tags=["spied"])
def _drive_mark_content_missing(session, record_id):
content_id = session.scalar(select(Asset.content_id).where(Asset.id == record_id))
mark_content_missing(session, content_id)
def _drive_update_asset_metadata(session, record_id):
update_asset_metadata(record_id, tags=["spied"])
def _drive_apply_tags(session, record_id):
apply_tags(record_id, ["spied"])
@pytest.mark.parametrize(
("target_module", "drive"),
[
(records_module, _drive_create_record),
(records_module, _drive_mark_content_missing),
(asset_management_module, _drive_update_asset_metadata),
(tagging_module, _drive_apply_tags),
],
ids=["create_record", "mark_content_missing", "update_asset_metadata", "apply_tags"],
)
def test_every_tag_write_site_routes_through_the_conflict_safe_helper(
session, mock_create_session, monkeypatch, target_module, drive
):
record_id = _seed_record(session, "/tmp/four-sites-fixture")
calls: list[str] = []
real_ensure_tag_link = records_module.ensure_tag_link
def spying_ensure_tag_link(*args, **kwargs):
calls.append(kwargs["tag_name"])
return real_ensure_tag_link(*args, **kwargs)
monkeypatch.setattr(target_module, "ensure_tag_link", spying_ensure_tag_link)
drive(session, record_id)
assert calls, (
"this site still writes asset_tags with a bare check-then-insert; a concurrent "
"writer inserting the same link makes it raise instead of settling"
)

View File

@@ -1,197 +1,52 @@
"""Tests for tagging services."""
import pytest
from sqlalchemy.orm import Session
from app.assets.database.models import Asset, AssetReference
from app.assets.database.queries import ensure_tags_exist, add_tags_to_reference
from app.assets.helpers import get_utc_now
from app.assets.services import apply_tags, remove_tags, list_tags
from app.assets.database.models import AssetTag, Tag
from app.assets.database.queries import create_content, create_record, fetch_record_tags
from app.assets.services.tagging import remove_tags
def _make_asset(session: Session, hash_val: str = "blake3:test") -> Asset:
asset = Asset(hash=hash_val, size_bytes=1024)
session.add(asset)
session.flush()
return asset
def test_tags_are_birth_facts_of_a_record(session: Session) -> None:
content = create_content(session, "/models/model.safetensors")
record = create_record(session, content.id, "model", tags=["model", "checkpoint"])
assert fetch_record_tags(session, record.id) == ["checkpoint", "model"]
def _make_reference(
session: Session,
asset: Asset,
name: str = "test",
owner_id: str = "",
) -> AssetReference:
now = get_utc_now()
ref = AssetReference(
owner_id=owner_id,
name=name,
asset_id=asset.id,
created_at=now,
updated_at=now,
last_access_time=now,
def _attach_automatic_tag(session: Session, record_id: str, tag_name: str) -> None:
if session.get(Tag, tag_name) is None:
session.add(Tag(name=tag_name))
session.flush()
session.add(
AssetTag(asset_id=record_id, tag_name=tag_name, origin="automatic")
)
session.add(ref)
session.flush()
return ref
class TestApplyTags:
def test_adds_new_tags(self, mock_create_session, session: Session):
asset = _make_asset(session)
ref = _make_reference(session, asset)
session.commit()
def test_removing_present_automatic_tag_reports_it_protected_not_absent(
session: Session, mock_create_session
) -> None:
content = create_content(session, "/output/protected.png")
record = create_record(session, content.id, "protected-fixture")
_attach_automatic_tag(session, record.id, "auto")
session.commit()
result = apply_tags(
reference_id=ref.id,
tags=["alpha", "beta"],
)
result = remove_tags(record.id, ["auto"])
assert set(result.added) == {"alpha", "beta"}
assert result.already_present == []
assert set(result.total_tags) == {"alpha", "beta"}
def test_reports_already_present(self, mock_create_session, session: Session):
asset = _make_asset(session)
ref = _make_reference(session, asset)
ensure_tags_exist(session, ["existing"])
add_tags_to_reference(session, reference_id=ref.id, tags=["existing"])
session.commit()
result = apply_tags(
reference_id=ref.id,
tags=["existing", "new"],
)
assert result.added == ["new"]
assert result.already_present == ["existing"]
def test_raises_for_nonexistent_ref(self, mock_create_session):
with pytest.raises(ValueError, match="not found"):
apply_tags(reference_id="nonexistent", tags=["x"])
def test_raises_for_wrong_owner(self, mock_create_session, session: Session):
asset = _make_asset(session)
ref = _make_reference(session, asset, owner_id="user1")
session.commit()
with pytest.raises(PermissionError, match="not owner"):
apply_tags(
reference_id=ref.id,
tags=["new"],
owner_id="user2",
)
assert result.protected == ["auto"]
assert result.not_present == []
assert result.removed == []
class TestRemoveTags:
def test_removes_tags(self, mock_create_session, session: Session):
asset = _make_asset(session)
ref = _make_reference(session, asset)
ensure_tags_exist(session, ["a", "b", "c"])
add_tags_to_reference(session, reference_id=ref.id, tags=["a", "b", "c"])
session.commit()
def test_remove_tags_separates_removed_protected_and_absent(
session: Session, mock_create_session
) -> None:
content = create_content(session, "/output/mixed.png")
record = create_record(session, content.id, "mixed-fixture", tags=["keep-me"])
_attach_automatic_tag(session, record.id, "auto")
session.commit()
result = remove_tags(
reference_id=ref.id,
tags=["a", "b"],
)
result = remove_tags(record.id, ["keep-me", "auto", "ghost"])
assert set(result.removed) == {"a", "b"}
assert result.not_present == []
assert result.total_tags == ["c"]
def test_reports_not_present(self, mock_create_session, session: Session):
asset = _make_asset(session)
ref = _make_reference(session, asset)
ensure_tags_exist(session, ["present"])
add_tags_to_reference(session, reference_id=ref.id, tags=["present"])
session.commit()
result = remove_tags(
reference_id=ref.id,
tags=["present", "absent"],
)
assert result.removed == ["present"]
assert result.not_present == ["absent"]
def test_raises_for_nonexistent_ref(self, mock_create_session):
with pytest.raises(ValueError, match="not found"):
remove_tags(reference_id="nonexistent", tags=["x"])
def test_raises_for_wrong_owner(self, mock_create_session, session: Session):
asset = _make_asset(session)
ref = _make_reference(session, asset, owner_id="user1")
session.commit()
with pytest.raises(PermissionError, match="not owner"):
remove_tags(
reference_id=ref.id,
tags=["x"],
owner_id="user2",
)
class TestListTags:
def test_returns_tags_with_counts(self, mock_create_session, session: Session):
ensure_tags_exist(session, ["used", "unused"])
asset = _make_asset(session)
ref = _make_reference(session, asset)
add_tags_to_reference(session, reference_id=ref.id, tags=["used"])
session.commit()
rows, total = list_tags()
tag_dict = {name: count for name, count in rows}
assert tag_dict["used"] == 1
assert tag_dict["unused"] == 0
assert total == 2
def test_excludes_zero_counts(self, mock_create_session, session: Session):
ensure_tags_exist(session, ["used", "unused"])
asset = _make_asset(session)
ref = _make_reference(session, asset)
add_tags_to_reference(session, reference_id=ref.id, tags=["used"])
session.commit()
rows, total = list_tags(include_zero=False)
tag_names = {name for name, _ in rows}
assert "used" in tag_names
assert "unused" not in tag_names
def test_prefix_filter(self, mock_create_session, session: Session):
ensure_tags_exist(session, ["alpha", "beta", "alphabet"])
session.commit()
rows, _ = list_tags(prefix="alph")
tag_names = {name for name, _ in rows}
assert tag_names == {"alpha", "alphabet"}
def test_order_by_name(self, mock_create_session, session: Session):
ensure_tags_exist(session, ["zebra", "alpha", "middle"])
session.commit()
rows, _ = list_tags(order="name_asc")
names = [name for name, _ in rows]
assert names == ["alpha", "middle", "zebra"]
def test_pagination(self, mock_create_session, session: Session):
ensure_tags_exist(session, ["a", "b", "c", "d", "e"])
session.commit()
rows, total = list_tags(limit=2, offset=1, order="name_asc")
assert total == 5
assert len(rows) == 2
names = [name for name, _ in rows]
assert names == ["b", "c"]
def test_clamps_limit(self, mock_create_session, session: Session):
ensure_tags_exist(session, ["a"])
session.commit()
# Service should clamp limit to max 1000
rows, _ = list_tags(limit=2000)
assert len(rows) <= 1000
assert result.removed == ["keep-me"]
assert result.protected == ["auto"]
assert result.not_present == ["ghost"]
assert fetch_record_tags(session, record.id) == ["auto"]

Some files were not shown because too many files have changed in this diff Show More