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:7007d18582Merge-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:7007d18582Merge-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:7007d18582Merge-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:7007d18582Merge-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:eca2c74bffMerge-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:eca2c74bffMerge-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 (7efdd1d7excluded, 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:ebc2cfeebcMerge-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:
@@ -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")
|
||||
|
||||
@@ -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")
|
||||
|
||||
125
alembic_db/versions/0007_record_content_split.py
Normal file
125
alembic_db/versions/0007_record_content_split.py
Normal 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"])
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
@@ -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
|
||||
333
app/assets/database/queries/records.py
Normal file
333
app/assets/database/queries/records.py
Normal 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()
|
||||
@@ -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
342
app/assets/docs/README.md
Normal 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.
|
||||
@@ -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
152
app/assets/lifecycle.py
Normal 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
227
app/assets/manager.py
Normal 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
30
app/assets/mode.py
Normal 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)
|
||||
@@ -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
|
||||
|
||||
97
app/assets/scanner_admission.py
Normal file
97
app/assets/scanner_admission.py
Normal 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
|
||||
205
app/assets/scanner_changes.py
Normal file
205
app/assets/scanner_changes.py
Normal 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)),
|
||||
)
|
||||
)
|
||||
)
|
||||
@@ -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)
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
|
||||
@@ -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 = (
|
||||
|
||||
@@ -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)
|
||||
189
app/assets/services/hash_mode_state.py
Normal file
189
app/assets/services/hash_mode_state.py
Normal 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
96
app/assets/services/lookup.py
Normal file
96
app/assets/services/lookup.py
Normal 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)
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
51
app/assets/services/snapshot_hash.py
Normal file
51
app/assets/services/snapshot_hash.py
Normal 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
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
11
comfy_execution/server_protocol.py
Normal file
11
comfy_execution/server_protocol.py
Normal 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: ...
|
||||
54
execution.py
54
execution.py
@@ -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
94
main.py
@@ -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()
|
||||
|
||||
18
openapi.yaml
18
openapi.yaml
@@ -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
|
||||
|
||||
60
server.py
60
server.py
@@ -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:
|
||||
|
||||
155
tests-unit/app_test/test_db_init_locking.py
Normal file
155
tests-unit/app_test/test_db_init_locking.py
Normal 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
|
||||
215
tests-unit/app_test/test_migration_0007.py
Normal file
215
tests-unit/app_test/test_migration_0007.py
Normal 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}"
|
||||
56
tests-unit/app_test/test_migration_roundtrip.py
Normal file
56
tests-unit/app_test/test_migration_roundtrip.py
Normal 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
|
||||
110
tests-unit/app_test/test_migration_warning.py
Normal file
110
tests-unit/app_test/test_migration_warning.py
Normal 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"
|
||||
)
|
||||
0
tests-unit/assets_test/__init__.py
Normal file
0
tests-unit/assets_test/__init__.py
Normal 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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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."""
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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) == []
|
||||
|
||||
171
tests-unit/assets_test/queries/test_lookup.py
Normal file
171
tests-unit/assets_test/queries/test_lookup.py
Normal 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"
|
||||
)
|
||||
@@ -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"}
|
||||
|
||||
137
tests-unit/assets_test/queries/test_records.py
Normal file
137
tests-unit/assets_test/queries/test_records.py
Normal 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"
|
||||
)
|
||||
@@ -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) == []
|
||||
|
||||
198
tests-unit/assets_test/queries/test_tags_query.py
Normal file
198
tests-unit/assets_test/queries/test_tags_query.py
Normal 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
|
||||
@@ -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:
|
||||
|
||||
35
tests-unit/assets_test/services/path_prefix_cases.py
Normal file
35
tests-unit/assets_test/services/path_prefix_cases.py
Normal 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}
|
||||
170
tests-unit/assets_test/services/test_admission_gate.py
Normal file
170
tests-unit/assets_test/services/test_admission_gate.py
Normal 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 == []
|
||||
@@ -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))) == []
|
||||
155
tests-unit/assets_test/services/test_api_routes_b.py
Normal file
155
tests-unit/assets_test/services/test_api_routes_b.py
Normal 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"
|
||||
@@ -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
|
||||
|
||||
@@ -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"
|
||||
)
|
||||
|
||||
@@ -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"
|
||||
)
|
||||
|
||||
@@ -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
|
||||
|
||||
113
tests-unit/assets_test/services/test_cached_execution_path.py
Normal file
113
tests-unit/assets_test/services/test_cached_execution_path.py
Normal 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)
|
||||
75
tests-unit/assets_test/services/test_cached_save.py
Normal file
75
tests-unit/assets_test/services/test_cached_save.py
Normal 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)
|
||||
115
tests-unit/assets_test/services/test_delete_b.py
Normal file
115
tests-unit/assets_test/services/test_delete_b.py
Normal 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
|
||||
197
tests-unit/assets_test/services/test_detection_gate.py
Normal file
197
tests-unit/assets_test/services/test_detection_gate.py
Normal 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
|
||||
@@ -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
|
||||
|
||||
308
tests-unit/assets_test/services/test_enrichment_predicate.py
Normal file
308
tests-unit/assets_test/services/test_enrichment_predicate.py
Normal 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)
|
||||
211
tests-unit/assets_test/services/test_enrichment_snapshot.py
Normal file
211
tests-unit/assets_test/services/test_enrichment_snapshot.py
Normal 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,
|
||||
}
|
||||
118
tests-unit/assets_test/services/test_from_hash.py
Normal file
118
tests-unit/assets_test/services/test_from_hash.py
Normal 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
|
||||
58
tests-unit/assets_test/services/test_hash_mode_state.py
Normal file
58
tests-unit/assets_test/services/test_hash_mode_state.py
Normal 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"
|
||||
@@ -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"
|
||||
|
||||
82
tests-unit/assets_test/services/test_ingest_b.py
Normal file
82
tests-unit/assets_test/services/test_ingest_b.py
Normal 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)
|
||||
168
tests-unit/assets_test/services/test_ingest_orphan_content.py
Normal file
168
tests-unit/assets_test/services/test_ingest_orphan_content.py
Normal 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"
|
||||
)
|
||||
333
tests-unit/assets_test/services/test_lifecycle.py
Normal file
333
tests-unit/assets_test/services/test_lifecycle.py
Normal 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}
|
||||
295
tests-unit/assets_test/services/test_manager_isolated.py
Normal file
295
tests-unit/assets_test/services/test_manager_isolated.py
Normal 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
|
||||
59
tests-unit/assets_test/services/test_metadata_update_b.py
Normal file
59
tests-unit/assets_test/services/test_metadata_update_b.py
Normal 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
|
||||
202
tests-unit/assets_test/services/test_noassets_parity.py
Normal file
202
tests-unit/assets_test/services/test_noassets_parity.py
Normal 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)
|
||||
198
tests-unit/assets_test/services/test_null_hash_recovery.py
Normal file
198
tests-unit/assets_test/services/test_null_hash_recovery.py
Normal 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
|
||||
143
tests-unit/assets_test/services/test_output_registration.py
Normal file
143
tests-unit/assets_test/services/test_output_registration.py
Normal 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)
|
||||
15
tests-unit/assets_test/services/test_owner_purge.py
Normal file
15
tests-unit/assets_test/services/test_owner_purge.py
Normal 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 == ""
|
||||
@@ -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
|
||||
168
tests-unit/assets_test/services/test_recovery_gate.py
Normal file
168
tests-unit/assets_test/services/test_recovery_gate.py
Normal 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
|
||||
273
tests-unit/assets_test/services/test_registration_primitives.py
Normal file
273
tests-unit/assets_test/services/test_registration_primitives.py
Normal 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)
|
||||
43
tests-unit/assets_test/services/test_scan_lifecycle.py
Normal file
43
tests-unit/assets_test/services/test_scan_lifecycle.py
Normal 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
|
||||
232
tests-unit/assets_test/services/test_scanner_b.py
Normal file
232
tests-unit/assets_test/services/test_scanner_b.py
Normal 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
|
||||
214
tests-unit/assets_test/services/test_scanner_seed_resilience.py
Normal file
214
tests-unit/assets_test/services/test_scanner_seed_resilience.py
Normal 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
|
||||
356
tests-unit/assets_test/services/test_serving_b.py
Normal file
356
tests-unit/assets_test/services/test_serving_b.py
Normal 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}"
|
||||
90
tests-unit/assets_test/services/test_snapshot_hash.py
Normal file
90
tests-unit/assets_test/services/test_snapshot_hash.py
Normal 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
|
||||
393
tests-unit/assets_test/services/test_split_policy.py
Normal file
393
tests-unit/assets_test/services/test_split_policy.py
Normal 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)
|
||||
282
tests-unit/assets_test/services/test_stored_hash_prefix.py
Normal file
282
tests-unit/assets_test/services/test_stored_hash_prefix.py
Normal 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)
|
||||
@@ -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
|
||||
|
||||
107
tests-unit/assets_test/services/test_tag_protection.py
Normal file
107
tests-unit/assets_test/services/test_tag_protection.py
Normal 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"]
|
||||
166
tests-unit/assets_test/services/test_tag_write_conflicts.py
Normal file
166
tests-unit/assets_test/services/test_tag_write_conflicts.py
Normal 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"
|
||||
)
|
||||
@@ -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
Reference in New Issue
Block a user