Files
ComfyAssets-ComfyUI_PromptM…/tests/test_review_schema.py
T
Vito 3f11ea6168 feat(annotations): Needs review queue backend (#284)
* feat(review): add review queue schema with restore validation

Add the Needs review schema group (media_review_metadata,
media_review_items, media_review_decisions) beside the image annotation
group it cites. review_schema_error validates the catalog, columns,
foreign keys, row values and canonical 8 MiB evidence read-only;
migrate_review_state installs it in a SAVEPOINT after the annotation
group, and verify_database_file requires it under require_library_state.
A review group without a valid annotation group is refused. Decisions
null their pair and image navigation through indexes when those are
deleted. Old backups and annotation-era backups restore and gain the
group; browser fixtures clear the new tables on reset.

* feat(review): persist ambiguous sampler evidence from scans and new outputs

The output scan and the live monitor now queue every file the ingestion
resolver finds ambiguous for Needs review. The file is captured, its metadata
is read again from the capture and resolved with the default profile, and only
an ambiguous result is stored, bound to the capture's annotation byte record
and fenced by its storage epoch. Evidence holds one hashed projection per
selectable candidate; none, more than 16 or over 8 MiB is a limited item.
Recording is idempotent per bytes and evidence and never touches a decision.

Scan statuses, counters, outcome rows and legacy SSE events are unchanged; the
job scope gains review_items/review_errors and the log a Needs review line.
The monitor records after its fallback link, best effort.

* feat(review): explicit association decisions with compatible linking

Needs review gains its queue reads and decisions:

- GET /prompt_manager/review/items lists pending items oldest first with
  the exact pending total and an HMAC keyset cursor bound to the storage
  epoch and the review index revision.
- GET /prompt_manager/review/items/{id} adds the candidates (role texts,
  label, selectability), the dependency fingerprint and the digest.
- POST /prompt_manager/review/items/{id}/decision applies a candidate only
  for a capture of the item's bytes that the current file still hashes to
  and whose evidence is unchanged, through the ingester's compatible-only
  rules in the same transaction as the decision row; a conflict writes
  nothing. Leave unresolved records the decision and links nothing.

ingest_resolved_pair keeps its contract; its body moves into
_ingest_resolved_pair_in, which runs in the caller's transaction.

Items keep the file's recorded spelling so targets keep their case on
Windows, and a rescan skips the copy and hash of a file whose size and
mtime match the bytes an item already cites.

* fix(review): build the spelling validation test from native paths

On Windows the locator is case-folded and uses backslashes, so the fixed
POSIX pair /out/a.png no longer folded to itself and the new spelling
check refused the good case. The test now spells a mixed-case native path
and folds it with os.path.normcase, and also refuses a locator of another
file.

* fix(review): Apply sees the file under every spelling and rechecks only its choice

Apply now finds the legacy rows of the same physical file however its root
is reached. The roots are real paths while the monitor links under the
directory it watches, so a symlinked or junctioned root, subst or a mapped
drive used to hide the monitor's row and Apply added a second, differently
linked row. Same-name rows are resolved to their physical identity
(normcase(realpath), else samefile); the dependency fingerprint covers all
of them, two rows refuse the choice, one row is enriched under the
ingester's compatible-only rules, and a same-name row the request did not
classify is a changed association.

Apply re-verifies only what the decision depends on: the bytes are still
ambiguous and still offer the chosen candidate. A resolver change that
rewords the reason or reorders candidates leaves pending items applicable.
The resolver version is one constant in prompt_pair_resolver, shared by
ingestion provenance and Needs review.

The transaction's byte check receives the re-hashed digest (or the body's
for Leave unresolved), every Apply refusal branch is tested with zero
writes, and the root listing and realpath calls run on a worker thread.

* fix(review): validate stored review evidence once per startup

Startup used to run the full review validation twice (the early guard and
the install) and the annotation group's row checks inside each, so every
evidence document was parsed and re-encoded twice. The early guard is now
catalog-only (a review group without annotation tables), and the install
validates the group once, trusting the annotation group it has just
validated. The evidence predicate no longer pattern-scans every document in
SQL, because the streamed canonical pass proves it, and the codec is built
once.

Measured on a database with 40,000 items of 4 KiB evidence: startup open
went from 7.7 s to 2.5 s. A test pins one evidence pass per reopen.
Restore and verify_database_file keep the full validation.

* fix(review): keep one capture worker for interactive requests

Needs review recording from the scan and the monitor shared the capture
store's two workers with the Tag dialog and Apply, so a busy scan could
answer a user's preview or decision with "Too many previews are open".
Background captures now never take the last worker or capture slot, and
wait up to two seconds for room before counting the file as not added.
Interactive captures are admitted or refused at once, as before.

* fix(review): refresh the recorded spelling when a file is seen under a new one

A case-only rename on Windows or macOS keeps size and mtime_ns, so the
rescan skipped the file and the item kept the old spelling. Its target was
then formatted from a name the directory no longer lists, the capture
refused it, and the item stayed pending with no way to apply it. The skip
path and a duplicate insert now record the current spelling for every item
at the locator and move the review index revision; item revisions do not
change, so an open decision stays valid. Tested under case-folding
semantics on every platform.

* fix(review): final capture store close and recoverable stale cursors

Closing the capture store at shutdown now runs under the store lock and is
final: a monitor recording or scan batch still in flight gets a capture
error instead of creating a new temp directory and restore subscription
after the cleanup.

A well-formed queue cursor this process cannot verify (opened before a
ComfyUI restart, whose cursor secret is per process) now answers 409
review-conflict with the current total, like any other stale page, instead
of an uncoded 400. Malformed cursors stay 400.

* fix(review): make only the shutdown cleanup final

The previous commit made every close of the capture store final, which
also refused the explicit closes that swap in a fresh store. Only the host's
on_cleanup hook is final now; an observe after it answers a coded capture
error instead of an unhandled 500.

* refactor(admin): move Needs review recording into its own mixin

_record_scanned_review, _record_review_sync, record_review_file and
_log_review_items move unchanged from py/api/admin.py to
py/api/review_recording.py (ReviewRecordingMixin), keeping admin.py from
growing further past the size ceiling. No behaviour change.

* fix(review): record the review fix round and the skip's residual risks in TESTING.md

The Needs review paragraphs now state the rescan skip's residual risks (an
in-place rewrite that keeps size and mtime_ns is never re-queued; a deleted
or moved file leaves its pending item counted), the physical identity rule
for Apply, the candidate-only recheck and shared resolver version, the
spelling refresh, background capture admission, the final shutdown close,
the stale-cursor answer and the measured startup validation cost.

* fix(review): refuse monitor recording beside record-changing jobs

System S1 (#289) gives every record change a registry admission class. The
scan already records Needs review items under its own records job. The
monitor's recording lasts one file and is not a job, so it now checks the
records admission: while a records or guarded job or hold runs it records
nothing, logs that the file waits for the next scan, and answers deferred.
The fallback link stands and the next "new" scan reads the ambiguous file
again, because a fallback link is not a resolved one.

JobRegistry.blocker and active iterate a snapshot of the jobs, since the
monitor's worker thread now asks them too.

* test(backup): wait for the restore, not a 5 s timer, in the thread test

test_database_is_usable_from_every_thread_after_restore let its worker
wait 5 s for the restore. On a stalled Windows runner (py3.11, where the
previous restore test also took 25 s instead of 0.3 s) the restore took
longer, the wait timed out, and the worker read the old database before
the restore ran. The worker now waits for the event under a 120 s bound
that only a hung test reaches, reports a timeout as an error, and the
main thread releases it even when the restore raises. With a restore
delayed by 6 s the old test fails and the new one passes.

* fix(review): store non-finite evidence as null and count pending items

canonical_evidence now follows main's non-finite policy (#303): NaN and
±Infinity are written as null instead of raising, so a NaN-stamped graph
never turns a review item limited. The System Overview reports the review
queue's pending total in needs_review.pending, read in the same snapshot
as the Library counts.

* fix(review): show the pending count on the System Overview review row

The Overview's review row now reads "N ambiguous image links" with the
concept's line, from needs_review.pending. Its Review metadata action
stays aria-disabled until the Needs review page ships. The system specs
and their concept comparison follow the real count.

* fix(system): answer the Overview and Statistics only after every read ended

asyncio.gather answered at the first failed read while its sibling
executor reads were still inside SQLite. In the suite the test teardown
then deleted the database under the running read, and py3.10 on Ubuntu
segfaulted in get_statistics (test_statistics_failure_is_an_error_envelope).
Both routes now wait for all their reads before answering.

* fix(review): queue only ambiguous files that nothing links, and heal the queue

The maintainer's live scan (19,908 files, scope "new") reported 4,677 files
under Needs review, but the queue held 10,232 items: every file the strict
resolver found ambiguous was queued, including the ones the legacy
browser-parser fallback then linked (all sampled "added" rows) and the ones
already linked to a known prompt (135 of 200 sampled "known" rows).

A file is now queued only when its final outcome is ambiguous and unlinked:
the resolver finds it ambiguous, the fallback links nothing and no saved row
links it already. An ambiguous file that a saved row links is reported as
unchanged ("Already linked to its prompt"), not as Needs review.

The queue heals itself. After each batch, the scan removes the pending items
of every listed file (read, or skipped as known) that a saved row now links.
The monitor does the same for each ambiguous output it links, and never
queues one. Only an undecided item is removed (pending, revision 0, no
decision row), so a decision is never deleted or changed and the schema
needs no new state.

The scan job's scope.review_items now counts this scan's files that are in
the pending queue after it, new or already there, so it agrees with the
review group's ambiguous rows and with the queue total. review_added and
review_settled are new, and the log gains "already in Needs review" and
"left Needs review · now linked" lines. A duplicate whose item is decided is
reported as "decided" and is not counted.

* fix(review): never queue an ambiguous file with nothing to choose

The maintainer's live queue held 39 items whose candidates all read
"unresolved" ("No positive / No negative"). An ambiguous result is now
reviewable only when at least two candidates are resolved with non-blank
positive text. With none the file could not be traced; with exactly one,
the untraced branches may still have made the file, so that text is not
certain either. Both are reported unresolved, never linked and never
queued (also the A1111 two-delimiter case, which has no candidates).

A rescan removes the pending items of such files whose own evidence
offers no choice, with the same safe delete as the linked settle
(pending, revision 0, no decision row). A limited item (over 16 candidates) that holds at least two
candidates is kept. The job scope gains review_untraced and the log
"= N images left Needs review · nothing to choose".

* test(scan): pin only the per-source scope keys beside the review counters

* test(system): match the scan scope keys beside the review counters
2026-10-23 14:19:11 -07:00

919 lines
38 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""Needs review schema group (#249): install, validation and restore safety."""
import hashlib
import json
import re
import sqlite3
import tempfile
import unittest
from contextlib import closing
from pathlib import Path
from unittest.mock import patch
try:
from tests.open_handles import assert_closed
except ImportError: # discovered with tests/ as the top-level directory
from open_handles import assert_closed
from database.models import PromptModel
from database.operations import PromptDatabase
try: # guarded so a missing module fails assertions, not test loading
from database.review_schema import (
canonical_evidence,
evidence_digest,
migrate_review_state,
review_schema_error,
)
except ImportError:
canonical_evidence = None
evidence_digest = None
migrate_review_state = None
review_schema_error = None
METADATA = "media_review_metadata"
ITEMS = "media_review_items"
DECISIONS = "media_review_decisions"
TABLES = (METADATA, ITEMS, DECISIONS)
ANNOTATION_TABLES = (
"media_annotation_metadata",
"media_annotation_files",
"media_annotation_bytes",
"media_image_tags",
)
MAX_EVIDENCE_BYTES = 8 * 1024 * 1024
# The exact DDL the plan pins; the module must produce precisely these objects.
DDL = {
METADATA: (
f"CREATE TABLE {METADATA}(singleton_id INTEGER PRIMARY KEY "
"CHECK(singleton_id=1), schema_version INTEGER NOT NULL "
"CHECK(schema_version=1), review_index_revision INTEGER NOT NULL "
"DEFAULT 0 CHECK(review_index_revision>=0))"
),
ITEMS: (
f"CREATE TABLE {ITEMS}(id TEXT PRIMARY KEY CHECK(length(id)=36), "
"byte_id TEXT NOT NULL REFERENCES media_annotation_bytes(id) "
"ON DELETE RESTRICT, spelled_path TEXT NOT NULL CHECK(length(CAST("
"spelled_path AS BLOB)) BETWEEN 1 AND 16384), "
"resolver_version INTEGER NOT NULL, "
"evidence_hash TEXT NOT NULL CHECK(length(evidence_hash)=64 AND "
"evidence_hash NOT GLOB '*[^0-9a-f]*'), evidence TEXT NOT NULL "
f"CHECK(length(CAST(evidence AS BLOB)) <= {MAX_EVIDENCE_BYTES}), "
"candidate_count INTEGER NOT NULL CHECK(candidate_count BETWEEN 0 AND 16), "
"limited INTEGER NOT NULL CHECK(limited IN (0,1)), state TEXT NOT NULL "
"CHECK(state IN ('pending','selected','unresolved')), revision INTEGER "
"NOT NULL CHECK(revision>=0), created_at TEXT NOT NULL, updated_at TEXT "
"NOT NULL, UNIQUE(byte_id, resolver_version, evidence_hash))"
),
DECISIONS: (
f"CREATE TABLE {DECISIONS}(id TEXT PRIMARY KEY CHECK(length(id)=36), "
f"item_id TEXT NOT NULL REFERENCES {ITEMS}(id) ON DELETE RESTRICT, "
"decision_revision INTEGER NOT NULL, created_at TEXT NOT NULL, "
"outcome TEXT NOT NULL CHECK(outcome IN ('selected','unresolved')), "
"candidate_hash TEXT CHECK(candidate_hash IS NULL OR "
"(length(candidate_hash)=64 AND candidate_hash NOT GLOB '*[^0-9a-f]*')), "
"evidence TEXT NOT NULL "
f"CHECK(length(CAST(evidence AS BLOB)) <= {MAX_EVIDENCE_BYTES}), "
"pair_id INTEGER REFERENCES prompt_pairs(id) ON DELETE SET NULL, "
"image_id INTEGER REFERENCES generated_images(id) ON DELETE SET NULL, "
"UNIQUE(item_id, decision_revision), "
"CHECK((outcome='selected') = (candidate_hash IS NOT NULL)))"
),
}
INDEXES = {
"idx_media_review_items_queue": (
f"CREATE INDEX idx_media_review_items_queue ON {ITEMS}"
"(state, created_at, id)"
),
"idx_media_review_decisions_pair": (
f"CREATE INDEX idx_media_review_decisions_pair ON {DECISIONS}(pair_id)"
),
"idx_media_review_decisions_image": (
f"CREATE INDEX idx_media_review_decisions_image ON {DECISIONS}(image_id)"
),
}
AUTOINDEXES = {ITEMS: 2, DECISIONS: 2}
LEGACY_TABLES = (
"prompts",
"generated_images",
"tags",
"prompt_tags",
"prompt_pairs",
"prompt_pair_tags",
"image_favorites",
) + ANNOTATION_TABLES
FILE_ID = "0f8e4b2a-1c3d-4e5f-8a9b-0c1d2e3f4a5b"
BYTE_ID = "6a1b2c3d-4e5f-4a6b-8c9d-0e1f2a3b4c5d"
ITEM_ID = "1d2e3f4a-5b6c-4d7e-8f9a-0b1c2d3e4f5a"
DECISION_ID = "2e3f4a5b-6c7d-4e8f-9a0b-1c2d3e4f5a6b"
OTHER_ID = "7b2c3d4e-5f6a-4b7c-9d0e-1f2a3b4c5d6e"
DIGEST = hashlib.sha256(b"img").hexdigest()
CANDIDATE = hashlib.sha256(b"candidate").hexdigest()
NOW = "2026-10-22T00:00:00.000000+00:00"
EVIDENCE = '{"candidates":[],"sha256":"' + DIGEST + '"}'
EVIDENCE_HASH = hashlib.sha256(EVIDENCE.encode("ascii")).hexdigest()
def _compact(sql):
return re.sub(r"\s+", "", sql or "").lower()
def expected_catalog():
objects = {("table", name, _compact(sql)) for name, sql in DDL.items()}
objects |= {("index", name, _compact(sql)) for name, sql in INDEXES.items()}
objects |= {
("index", f"sqlite_autoindex_{table}_{n}", "")
for table, count in AUTOINDEXES.items()
for n in range(1, count + 1)
}
return objects
def catalog(conn):
names = tuple(TABLES) + tuple(INDEXES)
marks = ",".join("?" * len(names))
table_marks = ",".join("?" * len(TABLES))
return {
(kind, name, _compact(sql))
for kind, name, sql in conn.execute(
"SELECT type,name,sql FROM sqlite_master "
f"WHERE name IN ({marks}) OR tbl_name IN ({table_marks})", # nosec B608
names + tuple(TABLES),
)
}
def rows(conn, sql, *params):
"""Plain tuples whatever row factory *conn* uses."""
return [tuple(row) for row in conn.execute(sql, params)]
def metadata(conn):
return rows(conn, f"SELECT schema_version,review_index_revision FROM {METADATA}")
def table_rows(conn, tables):
found = {}
for table in tables:
present = conn.execute(
"SELECT 1 FROM sqlite_master WHERE type='table' AND name=?", (table,)
).fetchone()
if present:
found[table] = rows(conn, f"SELECT * FROM {table}") # nosec B608
return found
def insert_byte(conn, file_id=FILE_ID, byte_id=BYTE_ID, locator="/out/img.png"):
conn.execute(
"INSERT INTO media_annotation_files VALUES(?,?,?)", (file_id, locator, NOW)
)
conn.execute(
"INSERT INTO media_annotation_bytes VALUES(?,?,?,?,?,?,?,0)",
(byte_id, file_id, DIGEST, 3, 1_700_000_000_000_000_000, NOW, NOW),
)
def insert_item(conn, *, item_id=ITEM_ID, evidence=EVIDENCE, candidates=2, **extra):
values = {
"id": item_id,
"byte_id": BYTE_ID,
"spelled_path": "/out/img.png",
"resolver_version": 1,
"evidence_hash": hashlib.sha256(evidence.encode("utf-8")).hexdigest(),
"evidence": evidence,
"candidate_count": candidates,
"limited": 0,
"state": "pending",
"revision": 0,
"created_at": NOW,
"updated_at": NOW,
}
values.update(extra)
columns = ",".join(values)
marks = ",".join("?" * len(values))
conn.execute(
f"INSERT INTO {ITEMS}({columns}) VALUES({marks})", # nosec B608
tuple(values.values()),
)
def insert_decision(conn, *, decision_id=DECISION_ID, **extra):
values = {
"id": decision_id,
"item_id": ITEM_ID,
"decision_revision": 1,
"created_at": NOW,
"outcome": "selected",
"candidate_hash": CANDIDATE,
"evidence": EVIDENCE,
"pair_id": None,
"image_id": None,
}
values.update(extra)
columns = ",".join(values)
marks = ",".join("?" * len(values))
conn.execute(
f"INSERT INTO {DECISIONS}({columns}) VALUES({marks})", # nosec B608
tuple(values.values()),
)
def insert_group_rows(conn, **item):
insert_byte(conn)
insert_item(conn, **item)
insert_decision(conn)
def populate_legacy(conn):
conn.execute(
"INSERT INTO prompts(id,text,hash,category,notes,run_count) VALUES"
"(7,'positive','p','keep','notes',12)"
)
conn.execute("INSERT INTO prompts(id,text,hash) VALUES(9,'negative','n')")
conn.execute(
"INSERT INTO generated_images(id,prompt_id,image_path,filename,file_path) "
"VALUES(4,7,'/tmp/image.png','image.png','/tmp/image.png')"
)
conn.execute("INSERT INTO tags(id,name) VALUES(3,'keep-tag')")
conn.execute("INSERT INTO prompt_tags VALUES(7,3)")
conn.execute(
"INSERT INTO prompt_pairs(id,positive_prompt_id,negative_prompt_id) "
"VALUES(6,7,9)"
)
conn.execute("INSERT INTO prompt_pair_tags VALUES(6,3)")
conn.execute("INSERT INTO image_favorites VALUES('root','a.png',?)", (NOW,))
def full_capacity_evidence():
"""Sixteen candidates of control, lone-surrogate, astral and combining
text, exactly 8 MiB."""
text = '\x00\x07\x1f
\ud800 é é "q" \\ \U0001f9ea'
candidates = [
{
"hash": hashlib.sha256(str(n).encode()).hexdigest(),
"negative": {"text": text * 3},
"positive": {"text": text * 5},
"sampler_id": str(n),
"state": "resolved",
}
for n in range(16)
]
value = {"candidates": candidates, "pad": ""}
base = len(canonical_evidence(value).encode("utf-8"))
value["pad"] = "x" * (MAX_EVIDENCE_BYTES - base)
return value
class ReviewSchemaTests(unittest.TestCase):
def setUp(self):
self.temp = tempfile.TemporaryDirectory()
self.addCleanup(self.temp.cleanup)
self.addCleanup(assert_closed, self, self.temp.name)
self.addCleanup(PromptModel.reset_schema_cache)
self.models = []
self.addCleanup(self.close_models)
self.path = str(Path(self.temp.name) / "live.db")
def close_models(self):
for model in self.models:
model.close_all()
def model(self, path=None):
model = PromptModel(path or self.path)
self.models.append(model)
return model
def require_module(self):
self.assertIsNotNone(
review_schema_error, "database.review_schema is not implemented"
)
def drop_tables(self, path, tables=TABLES):
"""Remove tables from a closed database file (an old backup)."""
with closing(sqlite3.connect(path)) as raw, raw:
raw.execute("PRAGMA foreign_keys=OFF")
for table in reversed(tables):
raw.execute(f"DROP TABLE {table}")
def test_fresh_install_and_reopen_is_idempotent(self):
self.require_module()
model = self.model()
with model.get_connection() as conn:
self.assertEqual(catalog(conn), expected_catalog())
self.assertIsNone(review_schema_error(conn, allow_missing=False))
self.assertEqual(metadata(conn), [(1, 0)])
insert_group_rows(conn)
snapshot = table_rows(conn, TABLES)
self.assertEqual(
PromptModel.verify_database_file(self.path, require_library_state=True),
(True, "ok"),
)
model.close_all()
for _ in range(2):
PromptModel.reset_schema_cache()
model = self.model()
with model.get_connection() as conn:
self.assertEqual(catalog(conn), expected_catalog())
self.assertIsNone(review_schema_error(conn, allow_missing=False))
self.assertEqual(table_rows(conn, TABLES), snapshot)
# Running the migration again on a current database changes nothing.
migrate_review_state(conn)
self.assertEqual(table_rows(conn, TABLES), snapshot)
model.close_all()
# Items and decisions keep the byte record and the item they cite.
PromptModel.reset_schema_cache()
model = self.model()
with model.get_connection() as conn:
with self.assertRaises(sqlite3.IntegrityError):
conn.execute("DELETE FROM media_annotation_files")
with self.assertRaises(sqlite3.IntegrityError):
conn.execute(f"DELETE FROM {ITEMS}")
self.assertEqual(table_rows(conn, TABLES), snapshot)
def test_startup_validates_stored_evidence_once(self):
"""A reopen streams the evidence once and the annotation group twice
(its own install), never again for the review group (I3)."""
self.require_module()
import database.annotation_schema as annotation_module
import database.review_schema as review_module
model = self.model()
with model.get_connection() as conn:
insert_group_rows(conn)
model.close_all()
for _ in range(2):
PromptModel.reset_schema_cache()
with (
patch.object(
review_module,
"_evidence_error",
wraps=review_module._evidence_error,
) as evidence,
patch.object(
review_module,
"annotation_schema_error",
wraps=review_module.annotation_schema_error,
) as nested,
patch.object(
annotation_module,
"_value_error",
wraps=annotation_module._value_error,
) as annotation_rows,
):
model = self.model()
model.get_connection()
self.assertEqual(evidence.call_count, 1)
self.assertEqual(nested.call_count, 0)
self.assertEqual(annotation_rows.call_count, 2)
model.close_all()
def test_old_database_without_group_migrates_without_touching_legacy_rows(self):
self.require_module()
for label, dropped in (
("annotations only", TABLES),
("predates annotations", ANNOTATION_TABLES + TABLES),
):
with self.subTest(label):
path = str(Path(self.temp.name) / f"{label.replace(' ', '-')}.db")
model = self.model(path)
with model.get_connection() as conn:
populate_legacy(conn)
if dropped == TABLES:
insert_byte(conn)
model.close_all()
self.drop_tables(path, dropped)
with closing(sqlite3.connect(path)) as raw:
self.assertIsNone(review_schema_error(raw))
self.assertEqual(
review_schema_error(raw, allow_missing=False),
"Missing review queue schema",
)
legacy_rows_before = table_rows(raw, LEGACY_TABLES)
self.assertTrue(PromptModel.verify_database_file(path)[0])
self.assertFalse(
PromptModel.verify_database_file(path, require_library_state=True)[
0
]
)
PromptModel.reset_schema_cache()
model = self.model(path)
with model.get_connection() as conn:
self.assertEqual(catalog(conn), expected_catalog())
self.assertIsNone(review_schema_error(conn, allow_missing=False))
self.assertEqual(metadata(conn), [(1, 0)])
counts = [
conn.execute(f"SELECT COUNT(*) FROM {t}").fetchone()[0]
for t in (ITEMS, DECISIONS)
]
legacy_rows_after = {
table: value
for table, value in table_rows(conn, LEGACY_TABLES).items()
if table in legacy_rows_before
}
self.assertEqual(counts, [0, 0])
self.assertEqual(legacy_rows_after, legacy_rows_before)
self.assertEqual(
PromptModel.verify_database_file(path, require_library_state=True),
(True, "ok"),
)
model.close_all()
def test_annotations_only_installation_is_valid(self):
self.require_module()
model = self.model()
with model.get_connection() as conn:
insert_byte(conn)
model.close_all()
self.drop_tables(self.path)
with closing(sqlite3.connect(self.path)) as raw:
self.assertIsNone(review_schema_error(raw))
self.assertEqual(PromptModel.verify_database_file(self.path), (True, "ok"))
self.assertEqual(
PromptModel.verify_database_file(self.path, require_library_state=True),
(False, "Missing review queue schema"),
)
def test_review_group_without_annotation_group_is_invalid(self):
self.require_module()
model = self.model()
model.close_all()
self.drop_tables(self.path, ANNOTATION_TABLES)
with closing(sqlite3.connect(self.path)) as raw:
for allow_missing in (True, False):
error = review_schema_error(raw, allow_missing=allow_missing)
self.assertIsNotNone(error)
self.assertIn("image annotation", error.lower())
self.assertFalse(PromptModel.verify_database_file(self.path)[0])
PromptModel.reset_schema_cache()
with self.assertRaises(sqlite3.DatabaseError):
self.model()
# An invalid annotation group invalidates the review group beside it.
path = self.tampered(
"bad annotations",
("UPDATE media_annotation_bytes SET sha256=upper(sha256)",),
checks_off=True,
)
with closing(sqlite3.connect(path)) as raw:
self.assertIn("image annotation", review_schema_error(raw).lower())
def tampered(self, name, statements, *, checks_off=False, rows_in=True):
"""A closed current database whose group was altered after the fact."""
path = str(Path(self.temp.name) / f"{name.replace(' ', '-')}.db")
model = self.model(path)
with model.get_connection() as conn:
populate_legacy(conn)
if rows_in:
insert_group_rows(conn)
model.close_all()
with closing(sqlite3.connect(path)) as raw, raw:
raw.execute("PRAGMA foreign_keys=OFF")
if checks_off:
raw.execute("PRAGMA ignore_check_constraints=ON")
for statement in statements:
if isinstance(statement, tuple):
raw.execute(*statement)
else:
raw.execute(statement)
return path
def assert_rejected_unrepaired(self, path):
with closing(sqlite3.connect(path)) as raw:
schema_before = raw.execute(
"SELECT type,name,tbl_name,sql FROM sqlite_master ORDER BY name"
).fetchall()
rows_before = table_rows(raw, TABLES)
for allow_missing in (False, True):
error = review_schema_error(raw, allow_missing=allow_missing)
self.assertIsNotNone(error)
self.assertIn("review", error.lower())
self.assertFalse(PromptModel.verify_database_file(path)[0])
self.assertFalse(
PromptModel.verify_database_file(path, require_library_state=True)[0]
)
PromptModel.reset_schema_cache()
with self.assertRaises(sqlite3.DatabaseError):
self.model(path)
PromptModel.reset_schema_cache()
with closing(sqlite3.connect(path)) as raw:
self.assertEqual(
raw.execute(
"SELECT type,name,tbl_name,sql FROM sqlite_master ORDER BY name"
).fetchall(),
schema_before,
)
self.assertEqual(table_rows(raw, TABLES), rows_before)
def test_partial_or_malformed_group_is_rejected_without_repair(self):
self.require_module()
cases = {
"missing metadata": (f"DROP TABLE {METADATA}",),
"missing items": (f"DROP TABLE {DECISIONS}", f"DROP TABLE {ITEMS}"),
"missing decisions": (f"DROP TABLE {DECISIONS}",),
"reserved view": (
f"DROP TABLE {DECISIONS}",
f"CREATE VIEW {DECISIONS} AS SELECT 'x' AS id",
),
"extra trigger": (
f"CREATE TRIGGER media_review_trigger AFTER INSERT ON {ITEMS} "
"BEGIN SELECT 1; END",
),
"extra index": (
f"CREATE INDEX idx_media_review_extra ON {ITEMS}(updated_at)",
),
"extra column": (
f"DROP TABLE {DECISIONS}",
DDL[DECISIONS].replace("evidence TEXT", "note TEXT, evidence TEXT"),
),
"wrong primary key": (
f"DROP TABLE {DECISIONS}",
DDL[DECISIONS].replace("id TEXT PRIMARY KEY", "id TEXT", 1),
),
"wrong foreign key": (
f"DROP TABLE {DECISIONS}",
DDL[DECISIONS].replace(
"generated_images(id) ON DELETE SET NULL",
"generated_images(id) ON DELETE CASCADE",
),
INDEXES["idx_media_review_decisions_pair"],
INDEXES["idx_media_review_decisions_image"],
),
"wrong check": (
f"DROP TABLE {METADATA}",
DDL[METADATA].replace("schema_version=1", "schema_version=2"),
f"INSERT INTO {METADATA} VALUES(1,2,0)",
),
"missing unique": (
"DROP INDEX idx_media_review_decisions_pair",
"DROP INDEX idx_media_review_decisions_image",
f"DROP TABLE {DECISIONS}",
DDL[DECISIONS].replace(" UNIQUE(item_id, decision_revision),", ""),
INDEXES["idx_media_review_decisions_pair"],
INDEXES["idx_media_review_decisions_image"],
),
"missing queue index": ("DROP INDEX idx_media_review_items_queue",),
"wrong index columns": (
"DROP INDEX idx_media_review_items_queue",
INDEXES["idx_media_review_items_queue"].replace(", id", ""),
),
"orphan item": ((f"UPDATE {ITEMS} SET byte_id=?", (OTHER_ID,)),),
"orphan decision": ((f"UPDATE {DECISIONS} SET item_id=?", (OTHER_ID,)),),
"dangling pair": (f"UPDATE {DECISIONS} SET pair_id=404",),
"bad item uuid": (
(f"UPDATE {ITEMS} SET id=?", ("x" * 36,)),
(f"UPDATE {DECISIONS} SET item_id=?", ("x" * 36,)),
),
"uppercase decision uuid": (f"UPDATE {DECISIONS} SET id=upper(id)",),
"null decision id": (f"UPDATE {DECISIONS} SET id=NULL",),
"blob ids": (f"UPDATE {DECISIONS} SET id=CAST(id AS BLOB)",),
"no metadata row": (f"DELETE FROM {METADATA}",),
"long created_at": ((f"UPDATE {ITEMS} SET created_at=?", ("x" * 65,)),),
"blob updated_at": (f"UPDATE {ITEMS} SET updated_at=X'00'",),
"text resolver version": (f"UPDATE {ITEMS} SET resolver_version='one'",),
"blob spelled path": (
f"UPDATE {ITEMS} SET spelled_path=CAST(spelled_path AS BLOB)",
),
"evidence hash mismatch": (
(f"UPDATE {ITEMS} SET evidence_hash=?", (CANDIDATE,)),
),
}
for label, statements in cases.items():
with self.subTest(label):
self.assert_rejected_unrepaired(self.tampered(label, statements))
unchecked = {
"two metadata rows": (f"INSERT INTO {METADATA} VALUES(2,1,0)",),
"wrong schema version": (f"UPDATE {METADATA} SET schema_version=2",),
"negative index revision": (
f"UPDATE {METADATA} SET review_index_revision=-1",
),
"uppercase evidence hash": (
f"UPDATE {ITEMS} SET evidence_hash=upper(evidence_hash)",
),
"bad state": (f"UPDATE {ITEMS} SET state='done'",),
"limited two": (f"UPDATE {ITEMS} SET limited=2",),
"empty spelled path": (f"UPDATE {ITEMS} SET spelled_path=''",),
"negative revision": (f"UPDATE {ITEMS} SET revision=-1",),
"negative candidates": (f"UPDATE {ITEMS} SET candidate_count=-1",),
"bad outcome": (f"UPDATE {DECISIONS} SET outcome='maybe'",),
"short candidate hash": (
f"UPDATE {DECISIONS} SET candidate_hash=substr(candidate_hash,1,63)",
),
"negative decision revision": (
f"UPDATE {DECISIONS} SET decision_revision=-1",
),
}
for label, statements in unchecked.items():
with self.subTest(label):
path = self.tampered(label, statements, checks_off=True)
self.assert_rejected_unrepaired(path)
def test_selected_decision_requires_candidate_hash(self):
self.require_module()
model = self.model()
with model.get_connection() as conn:
insert_byte(conn)
insert_item(conn)
for outcome, candidate in (("selected", None), ("unresolved", CANDIDATE)):
with self.subTest(outcome):
with self.assertRaises(sqlite3.IntegrityError):
insert_decision(conn, outcome=outcome, candidate_hash=candidate)
insert_decision(conn, outcome="unresolved", candidate_hash=None)
insert_decision(conn, decision_id=OTHER_ID, decision_revision=2)
self.assertIsNone(review_schema_error(conn, allow_missing=False))
self.assertEqual(
rows(conn, f"SELECT outcome FROM {DECISIONS} ORDER BY 1"),
[("selected",), ("unresolved",)],
)
with self.assertRaises(sqlite3.IntegrityError): # one per revision
insert_decision(conn, decision_id=ITEM_ID, decision_revision=2)
model.close_all()
path = self.tampered(
"selected without hash",
(f"UPDATE {DECISIONS} SET candidate_hash=NULL",),
checks_off=True,
)
self.assert_rejected_unrepaired(path)
def test_oversize_or_noncanonical_evidence_rejected(self):
self.require_module()
self.assertEqual(
canonical_evidence({"b": [1, "é"], "a": None}),
'{"a":null,"b":[1,"\\u00e9"]}',
)
self.assertEqual(
evidence_digest(EVIDENCE), hashlib.sha256(EVIDENCE.encode()).hexdigest()
)
# The store's non-finite policy (utils.json_strict): written as null.
self.assertEqual(
canonical_evidence({"x": float("nan"), "y": [float("-inf")]}),
'{"x":null,"y":[null]}',
)
for bad in (
[],
"text",
{"x": object()},
{"x": object(), "y": float("nan")},
{1: 1, "a": 2},
):
with self.subTest(value=repr(bad)), self.assertRaises(ValueError):
canonical_evidence(bad)
with self.assertRaises(ValueError):
canonical_evidence({"pad": "x" * MAX_EVIDENCE_BYTES})
oversize = '{"pad":"' + "x" * (MAX_EVIDENCE_BYTES - 9) + '"}'
self.assertEqual(len(oversize), MAX_EVIDENCE_BYTES + 1)
model = self.model()
with model.get_connection() as conn:
insert_byte(conn)
with self.assertRaises(sqlite3.IntegrityError):
insert_item(conn, evidence=oversize)
model.close_all()
noncanonical = {
"spaces": '{"candidates": [], "sha256": "x"}',
"unsorted": '{"sha256":"x","candidates":[]}',
"raw unicode": '{"text":"é"}',
"duplicate key": '{"a":1,"a":2}',
"not an object": "[]",
"not json": "{",
"nan": '{"a":NaN}',
"empty": "",
}
for label, evidence in noncanonical.items():
for table in (ITEMS, DECISIONS):
with self.subTest(label, table=table):
digest = hashlib.sha256(evidence.encode("utf-8")).hexdigest()
statements = [
(f"UPDATE {table} SET evidence=?", (evidence,)),
(f"UPDATE {ITEMS} SET evidence_hash=?", (digest,)),
]
if table == DECISIONS:
statements.pop()
self.assert_rejected_unrepaired(
self.tampered(f"{label} {table}", statements)
)
for table in (ITEMS, DECISIONS):
with self.subTest("oversize", table=table):
digest = hashlib.sha256(oversize.encode("ascii")).hexdigest()
statements = [
(f"UPDATE {table} SET evidence=?", (oversize,)),
(f"UPDATE {ITEMS} SET evidence_hash=?", (digest,)),
]
path = self.tampered(f"oversize {table}", statements, checks_off=True)
self.assert_rejected_unrepaired(path)
def test_seventeen_candidates_rejected(self):
self.require_module()
model = self.model()
with model.get_connection() as conn:
insert_byte(conn)
with self.assertRaises(sqlite3.IntegrityError):
insert_item(conn, candidates=17)
insert_item(conn, candidates=16)
self.assertIsNone(review_schema_error(conn, allow_missing=False))
model.close_all()
path = self.tampered(
"seventeen", (f"UPDATE {ITEMS} SET candidate_count=17",), checks_off=True
)
self.assert_rejected_unrepaired(path)
def test_prompt_or_image_deletion_nulls_navigation_and_keeps_decision(self):
self.require_module()
db = PromptDatabase(self.path)
self.models.append(db.model)
with db.model.get_connection() as conn:
populate_legacy(conn)
conn.execute(
"INSERT INTO prompt_pairs(id,positive_prompt_id,negative_prompt_id) "
"VALUES(8,9,NULL)"
)
insert_byte(conn)
insert_item(conn)
insert_decision(conn, pair_id=8, image_id=4)
conn.commit()
self.assertTrue(db.delete_prompt_pair(8))
with db.model.get_connection() as conn:
self.assertEqual(
rows(conn, f"SELECT pair_id, image_id FROM {DECISIONS}"), [(None, 4)]
)
with db.model.get_connection() as conn:
conn.execute("DELETE FROM prompt_pairs WHERE id=6")
conn.commit()
self.assertTrue(db.delete_prompt(7))
with db.model.get_connection() as conn:
self.assertEqual(
rows(
conn,
f"SELECT id, item_id, outcome, candidate_hash, pair_id, image_id "
f"FROM {DECISIONS}",
),
[(DECISION_ID, ITEM_ID, "selected", CANDIDATE, None, None)],
)
self.assertEqual(rows(conn, f"SELECT id FROM {ITEMS}"), [(ITEM_ID,)])
self.assertIsNone(review_schema_error(conn, allow_missing=False))
# The SET NULL actions look decisions up by index, not by a scan.
for column in ("pair_id", "image_id"):
plan = " ".join(
str(row[3])
for row in conn.execute(
f"EXPLAIN QUERY PLAN UPDATE {DECISIONS} " # nosec B608
f"SET {column}=NULL WHERE {column}=?",
(1,),
)
)
self.assertIn(f"idx_media_review_decisions_{column[:-3]}", plan)
def test_full_capacity_unicode_control_pair_round_trips_within_8_mib(self):
self.require_module()
value = full_capacity_evidence()
evidence = canonical_evidence(value)
self.assertEqual(len(evidence.encode("utf-8")), MAX_EVIDENCE_BYTES)
self.assertTrue(evidence.isascii())
source_path = str(Path(self.temp.name) / "source.db")
source = self.model(source_path)
with source.get_connection() as conn:
conn.execute("PRAGMA wal_autocheckpoint=0")
insert_byte(conn)
insert_item(conn, evidence=evidence, candidates=16)
insert_decision(conn, evidence=evidence)
conn.commit()
source_rows = table_rows(conn, TABLES)
self.assertGreater(Path(source_path + "-wal").stat().st_size, 0)
live = self.model()
backup_file = str(Path(self.temp.name) / "backup.db")
self.assertTrue(source.backup_database(backup_file))
for restore_from in (source_path, backup_file):
with self.subTest(Path(restore_from).name):
live.restore_from_file(restore_from)
with live.get_connection() as conn:
self.assertEqual(table_rows(conn, TABLES), source_rows)
self.assertIsNone(review_schema_error(conn, allow_missing=False))
stored = conn.execute(
f"SELECT i.evidence, d.evidence FROM {ITEMS} i "
f"JOIN {DECISIONS} d ON d.item_id = i.id"
).fetchone()
self.assertEqual(json.loads(stored[0]), value)
self.assertEqual(json.loads(stored[1]), value)
PromptModel.close_all_instances()
assert_closed(self, self.temp.name)
def test_install_failure_rolls_back_and_retries(self):
self.require_module()
with patch(
"database.review_schema.review_schema_error",
side_effect=[None, "injected final validation failure"],
):
with self.assertRaisesRegex(sqlite3.DatabaseError, "injected"):
self.model()
with closing(sqlite3.connect(self.path)) as raw:
self.assertEqual(catalog(raw), set())
self.assertIsNotNone(
raw.execute(
"SELECT 1 FROM sqlite_master WHERE name='media_annotation_bytes'"
).fetchone()
)
PromptModel.reset_schema_cache()
model = self.model()
self.assertEqual(
PromptModel.verify_database_file(model.db_path, require_library_state=True),
(True, "ok"),
)
with model.get_connection() as conn:
self.assertEqual(metadata(conn), [(1, 0)])
def test_wal_backup_restore_round_trips_group(self):
self.require_module()
source_path = str(Path(self.temp.name) / "source.db")
source = self.model(source_path)
with source.get_connection() as conn:
conn.execute("PRAGMA wal_autocheckpoint=0")
insert_group_rows(conn)
conn.execute(f"UPDATE {METADATA} SET review_index_revision=3")
conn.commit()
source_rows = table_rows(conn, TABLES)
self.assertGreater(Path(source_path + "-wal").stat().st_size, 0)
live = self.model()
with live.get_connection() as conn:
insert_byte(conn)
insert_item(conn, item_id=OTHER_ID)
source_before = Path(source_path).read_bytes()
backup_file = str(Path(self.temp.name) / "backup.db")
self.assertTrue(source.backup_database(backup_file))
with closing(sqlite3.connect(backup_file)) as raw:
self.assertEqual(table_rows(raw, TABLES), source_rows)
self.assertIsNone(review_schema_error(raw, allow_missing=False))
safety = live.restore_from_file(source_path)
self.assertEqual(Path(source_path).read_bytes(), source_before)
with live.get_connection() as conn:
self.assertEqual(table_rows(conn, TABLES), source_rows)
self.assertEqual(metadata(conn), [(1, 3)])
with closing(sqlite3.connect(safety)) as raw:
self.assertEqual(rows(raw, f"SELECT id FROM {ITEMS}"), [(OTHER_ID,)])
PromptModel.reset_schema_cache()
live.restore_from_file(backup_file)
with live.get_connection() as conn:
self.assertEqual(table_rows(conn, TABLES), source_rows)
PromptModel.close_all_instances()
assert_closed(self, self.temp.name)
def test_staged_failure_preserves_source_live_safety_and_handles(self):
self.require_module()
live = self.model()
with live.get_connection() as conn:
insert_byte(conn)
insert_item(conn, item_id=OTHER_ID)
live_rows = table_rows(conn, TABLES)
bad = self.tampered(
"restore bad", (f"UPDATE {ITEMS} SET state='done'",), checks_off=True
)
old = str(Path(self.temp.name) / "old.db")
self.model(old).close_all()
self.drop_tables(old)
for source, patched in ((bad, False), (old, True)):
with self.subTest(source=Path(source).name):
before = Path(source).read_bytes()
safety = sorted(Path(self.temp.name).glob("live.db.backup_*"))
if patched:
context = patch(
"database.models.migrate_review_state",
side_effect=sqlite3.OperationalError("staging failed"),
)
expected = sqlite3.OperationalError
else:
context = patch.object(PromptModel, "backup_database")
expected = ValueError
with context as hook, self.assertRaises(expected):
live.restore_from_file(source)
if not patched:
hook.assert_not_called()
self.assertEqual(Path(source).read_bytes(), before)
self.assertEqual(
sorted(Path(self.temp.name).glob("live.db.backup_*")), safety
)
assert_closed(self, source)
with live.get_connection() as conn:
self.assertEqual(table_rows(conn, TABLES), live_rows)
self.assertFalse(
PromptModel.verify_database_file(old, require_library_state=True)[0]
)
self.assertTrue(PromptModel.verify_database_file(old)[0])
def test_verify_requires_group_only_with_require_library_state(self):
self.require_module()
self.model().close_all()
self.assertEqual(
PromptModel.verify_database_file(self.path, require_library_state=True),
(True, "ok"),
)
self.drop_tables(self.path)
self.assertEqual(PromptModel.verify_database_file(self.path), (True, "ok"))
self.assertEqual(
PromptModel.verify_database_file(self.path, require_library_state=True),
(False, "Missing review queue schema"),
)
if __name__ == "__main__":
unittest.main()