Vito
d88cf9e7c6
* fix(quality): preserve cursor iteration and option validation Keep prefetched rows on ManagedCursor so partial loops and explicit fetches share one position, and stop existing iterators after close or retirement. Preserve chunked reads and native SQLite result lifetime on reexecution and errors. Catch OptionError before ValueError so legacy auto-tag requests receive the actual validation error. Remove the overwritten ViewerJS title option. Regression tests failed before the fixes for lost rows, reads after close/retirement, and misleading option errors. Focused Python checks pass; cursor checks also pass on Python 3.13 and 3.14. Refs #321. * fix(quality): make the test doubles legal raises and hashes The fake queue no longer raises an attribute that starts as None. The oversized kind records hashes instead of raising from __hash__, which also catches a hash that production code swallows. Hostile.__bool__ raises the TypeError a truth test may raise. * fix(quality): drop the extra saveCoverage argument * refactor(quality): remove assignments that are overwritten before use * refactor(quality): remove the remaining overwritten test assignments * refactor(quality): drop imports repeated inside functions * refactor(quality): give duplicate identity keys one shape * refactor(quality): name the unused job event loop variable * refactor(quality): release model gates in finally, not except BaseException * refactor(quality): end the job runners with an explicit return * refactor(quality): use the precise unittest asserts * refactor(quality): import unittest.mock one way in the tests * fix(quality): say why each empty except is safe, log the broad ones * refactor(quality): stop binding values the tests never read * refactor(quality): remove the unused preview aliases and tag page query * refactor(quality): remove re-exports nothing imports * refactor(quality): parenthesize deliberate string joins in lists * refactor(quality): remove unused JavaScript variables and imports * refactor(quality): use the precise assert in the model status test
462 lines
18 KiB
Python
462 lines
18 KiB
Python
"""Move library, final review round (#249 S6b PR #318).
|
|
|
|
- I1: a switch that fails at any point is undone while both files are still
|
|
locked: the copy is removed, the old library stays in use and a restart
|
|
keeps it. A follow-up after a complete switch only logs.
|
|
- I2: a background write (run counting, the image monitor) that fails
|
|
because Move library switched runs once more, on the new library; a write
|
|
racing the switch on a real connection (a run count, an image link) lands
|
|
in the new library only.
|
|
- The switching thread never closes a connection another thread is inside
|
|
a statement on (a use-after-free in SQLite: "Windows fatal exception:
|
|
access violation"); it interrupts it and the owner closes it.
|
|
- M2: the copy is verified through a file URI that survives ``#`` and ``%``.
|
|
- M6: a restart after a crash between the switch and the rename renames the
|
|
old files, never deletes them.
|
|
"""
|
|
|
|
import contextlib
|
|
import json
|
|
import os
|
|
import sqlite3
|
|
import sys
|
|
import threading
|
|
import unittest
|
|
import unittest.mock as mock
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
from database import library_move as lm # noqa: E402
|
|
from database import models # noqa: E402
|
|
from database.operations import PromptDatabase # noqa: E402
|
|
from tests.background_stops import close_handler # noqa: E402
|
|
from tests.open_handles import assert_closed # noqa: E402
|
|
from utils import library_location as ll # noqa: E402
|
|
from utils import usage_tracking # noqa: E402
|
|
from utils.hashing import generate_prompt_hash # noqa: E402
|
|
|
|
try:
|
|
from tests.test_library_move import KEY, SimulatedCrash, _MoveCase, _rows
|
|
except ImportError: # discovered with tests/ as the top-level directory
|
|
from test_library_move import KEY, SimulatedCrash, _MoveCase, _rows
|
|
|
|
WAIT = 30
|
|
|
|
|
|
def _run_count(path, text):
|
|
uri = "file:" + path.replace(os.sep, "/") + "?mode=ro"
|
|
with contextlib.closing(sqlite3.connect(uri, uri=True)) as conn:
|
|
row = conn.execute("SELECT run_count FROM prompts WHERE text=?", (text,))
|
|
return row.fetchone()[0]
|
|
|
|
|
|
class _CrossThreadCloseProbe:
|
|
"""Record every ``close()`` this thread calls on a connection in use elsewhere.
|
|
|
|
The writer marks its connection with ``using`` when its statement starts.
|
|
Closing that connection from the switching thread while the writer is
|
|
inside the statement frees memory SQLite is still using; on Windows the
|
|
process dies with an access violation, on Linux the writer sees garbage.
|
|
"""
|
|
|
|
def __init__(self):
|
|
self.in_use = set()
|
|
self.violations = []
|
|
|
|
def using(self, conn):
|
|
self.in_use.add(conn)
|
|
|
|
def _profile(self, _frame, event, arg):
|
|
if event != "c_call" or getattr(arg, "__name__", "") != "close":
|
|
return
|
|
target = getattr(arg, "__self__", None)
|
|
if isinstance(target, sqlite3.Connection) and target in self.in_use:
|
|
self.violations.append(threading.current_thread().name)
|
|
|
|
@contextlib.contextmanager
|
|
def watching(self):
|
|
previous = sys.getprofile()
|
|
sys.setprofile(self._profile)
|
|
try:
|
|
yield self
|
|
finally:
|
|
sys.setprofile(previous)
|
|
|
|
|
|
class TestFailedSwitchIsUndone(_MoveCase):
|
|
"""I1: whatever part of the switch fails, nothing was moved."""
|
|
|
|
def _assert_not_moved(self, caught):
|
|
self.assertEqual(caught.exception.code, "switch-failed")
|
|
self.assertIn("Nothing was moved", caught.exception.message)
|
|
self.assertIsNone(ll.current_override())
|
|
self.assertEqual(self.db.model.db_path, self.old_db)
|
|
# Writes keep landing in the old library, and a restart keeps it.
|
|
self.db.save_prompt("written after the failed switch")
|
|
self.assertIn("written after the failed switch", repr(_rows(self.old_db)))
|
|
self.assertEqual(self.location().source, ll.SOURCE_NODE_FOLDER)
|
|
self.assertNotEqual(ll.read_marker(self.target).get("state"), "moved")
|
|
self.assert_target_clean()
|
|
self.assertTrue(os.path.isfile(self.config))
|
|
assert_closed(self, self.target)
|
|
|
|
def _failing(self, point):
|
|
real_switch_to = ll.switch_to
|
|
real_repoint = models.PromptModel.repoint_instances
|
|
|
|
def switch(old_db, location, detached):
|
|
if point == "before":
|
|
raise RuntimeError("switch failed before anything")
|
|
if point == "resolver":
|
|
real_switch_to(location)
|
|
raise RuntimeError("switch failed after the resolver")
|
|
if point == "repoint":
|
|
real_switch_to(location)
|
|
detached.extend(real_repoint(old_db, location.db_path))
|
|
raise RuntimeError("switch failed after the models")
|
|
lm.switch_process(old_db, location, detached)
|
|
raise RuntimeError("switch failed at its very end")
|
|
|
|
return switch
|
|
|
|
def test_a_switch_failing_at_any_point_is_undone(self):
|
|
node = PromptDatabase(self.old_db)
|
|
node.get_prompt_by_id(1)
|
|
for point in ("before", "resolver", "repoint", "complete"):
|
|
with self.subTest(point=point):
|
|
with self.assertRaises(lm.MoveFailed) as caught:
|
|
self.move(switch=self._failing(point))
|
|
self._assert_not_moved(caught)
|
|
self.assertEqual(node.model.db_path, self.old_db)
|
|
self.assertEqual(node.get_prompt_by_id(1)["text"], "a cat number 0")
|
|
|
|
def test_a_failed_switch_closes_the_old_connections_it_detached(self):
|
|
node = PromptDatabase(self.old_db)
|
|
for point in ("repoint", "complete"):
|
|
with self.subTest(point=point):
|
|
before = node.model.get_connection()
|
|
with self.assertRaises(lm.MoveFailed):
|
|
self.move(switch=self._failing(point))
|
|
with self.assertRaises(sqlite3.ProgrammingError):
|
|
before.execute("SELECT 1")
|
|
# The model reopens the old library on next use.
|
|
self.assertEqual(node.get_prompt_by_id(1)["text"], "a cat number 0")
|
|
|
|
def test_a_failing_follow_up_after_the_switch_only_logs(self):
|
|
def follow_up():
|
|
raise RuntimeError("the health report could not be dropped")
|
|
|
|
with self.assertLogs("prompt_manager.database.library_move", "ERROR"):
|
|
result = self.move(after_switch=follow_up)
|
|
self.assertIsNone(result.retire_problem)
|
|
self.assertEqual(self.db.model.db_path, self.new_db)
|
|
self.assertTrue(
|
|
os.path.isfile(
|
|
os.path.join(self.ext, "prompts.db.moved-to-__promptmanager")
|
|
)
|
|
)
|
|
|
|
|
|
class _Flaky:
|
|
"""A write that fails once as if Move library switched while it waited."""
|
|
|
|
def __init__(self, db, name):
|
|
self.db, self.name, self.calls = db, name, 0
|
|
self.real = getattr(db, name)
|
|
|
|
def __call__(self, *args, **kwargs):
|
|
self.calls += 1
|
|
if self.calls == 1:
|
|
self.db.model.close_all() # what the switch does: a new generation
|
|
raise sqlite3.ProgrammingError("Cannot operate on a closed database.")
|
|
return self.real(*args, **kwargs)
|
|
|
|
|
|
class TestBackgroundWritesSurviveTheSwitch(_MoveCase):
|
|
"""I2: run counting and the image monitor write again on a fresh connection."""
|
|
|
|
def test_run_counting_retries_once_after_a_switch(self):
|
|
db = PromptDatabase(self.old_db)
|
|
text = "a counted prompt"
|
|
db.save_prompt(text, prompt_hash=generate_prompt_hash(text))
|
|
flaky = _Flaky(db, "record_prompt_use")
|
|
with mock.patch.object(db, "record_prompt_use", flaky):
|
|
pending = usage_tracking.PendingFirstUse()
|
|
usage_tracking.count_runs([(generate_prompt_hash(text), text)], db, pending)
|
|
self.assertEqual(flaky.calls, 2)
|
|
self.assertEqual(_run_count(self.old_db, text), 1)
|
|
|
|
def test_a_failure_without_a_switch_is_not_retried(self):
|
|
db = PromptDatabase(self.old_db)
|
|
calls = []
|
|
|
|
def broken():
|
|
calls.append(1)
|
|
raise sqlite3.OperationalError("disk I/O error")
|
|
|
|
with self.assertRaises(sqlite3.OperationalError):
|
|
db.survive_library_switch(broken)
|
|
self.assertEqual(calls, [1])
|
|
|
|
def test_an_image_link_racing_the_switch_lands_in_the_new_library(self):
|
|
from utils.image_monitor import ImageGenerationHandler
|
|
|
|
self.db.model.close_all()
|
|
image = os.path.join(self.comfy, "output", "late.png")
|
|
writer_db = PromptDatabase(self.old_db)
|
|
handler = ImageGenerationHandler(writer_db, mock.MagicMock())
|
|
self.addCleanup(close_handler, self, handler)
|
|
started = threading.Event()
|
|
probe = _CrossThreadCloseProbe()
|
|
outcome = {}
|
|
|
|
def write():
|
|
conn = writer_db.model.get_connection()
|
|
|
|
def trace(sql):
|
|
if sql.lstrip().upper().startswith(("BEGIN", "INSERT")):
|
|
probe.using(conn)
|
|
started.set()
|
|
|
|
conn.set_trace_callback(trace)
|
|
try:
|
|
# No metadata: the fallback link, not the resolved pair.
|
|
handler.link_image_to_prompt(image, {"id": 1}, {})
|
|
outcome["ok"] = True
|
|
except Exception as error: # pragma: no cover - reported below
|
|
outcome["error"] = error
|
|
finally:
|
|
writer_db.model.close()
|
|
|
|
thread = threading.Thread(target=write)
|
|
|
|
def checkpoint(step):
|
|
if step == "locked":
|
|
# The old file is write-locked now: start the writer and
|
|
# wait until its first write is running (it then waits on us).
|
|
thread.start()
|
|
self.assertTrue(started.wait(WAIT))
|
|
|
|
with probe.watching():
|
|
result = self.move(checkpoint=checkpoint)
|
|
thread.join(WAIT)
|
|
self.assertFalse(thread.is_alive())
|
|
self.assertEqual(probe.violations, [])
|
|
self.assertEqual(outcome, {"ok": True})
|
|
self.assertIsNone(result.retire_problem)
|
|
moved = [
|
|
row
|
|
for row in _rows(self.new_db)["generated_images"]
|
|
if "late.png" in repr(row)
|
|
]
|
|
self.assertEqual(len(moved), 1)
|
|
renamed = os.path.join(self.ext, "prompts.db.moved-to-__promptmanager")
|
|
self.assertNotIn("late.png", repr(_rows(renamed).get("generated_images")))
|
|
|
|
def test_an_image_that_was_not_linked_is_not_reported_as_linked(self):
|
|
from utils.image_monitor import ImageGenerationHandler
|
|
|
|
db = PromptDatabase(self.old_db)
|
|
handler = ImageGenerationHandler(db, mock.MagicMock())
|
|
self.addCleanup(close_handler, self, handler)
|
|
image = os.path.join(self.comfy, "output", "orphan.png")
|
|
with mock.patch.object(handler, "logger") as log:
|
|
handler.link_image_to_prompt(image, {"id": 99999}, {})
|
|
messages = [str(call) for call in log.debug.call_args_list]
|
|
self.assertFalse([m for m in messages if "Successfully linked" in m])
|
|
self.assertNotIn("orphan.png", repr(_rows(self.old_db)["generated_images"]))
|
|
|
|
def test_a_write_racing_the_switch_lands_in_the_new_library(self):
|
|
text = "a racing prompt"
|
|
self.db.save_prompt(text, prompt_hash=generate_prompt_hash(text))
|
|
self.db.model.close_all()
|
|
writer_db = PromptDatabase(self.old_db)
|
|
thread, started, probe, outcome = self._racing_run_count(writer_db, text)
|
|
|
|
def checkpoint(step):
|
|
if step == "locked":
|
|
# The old file is write-locked now: start the writer and
|
|
# wait until its UPDATE is running (it then waits on us).
|
|
thread.start()
|
|
self.assertTrue(started.wait(WAIT))
|
|
|
|
with probe.watching():
|
|
result = self.move(checkpoint=checkpoint)
|
|
# The writer's old connection was closed before the old file was
|
|
# renamed: Windows refuses to rename an open file.
|
|
self.assertIsNone(result.retire_problem)
|
|
thread.join(WAIT)
|
|
self.assertFalse(thread.is_alive())
|
|
self.assertEqual(probe.violations, [])
|
|
self.assertEqual(outcome, {"ok": True})
|
|
self.assertEqual(_run_count(self.new_db, text), 1)
|
|
renamed = os.path.join(self.ext, "prompts.db.moved-to-__promptmanager")
|
|
self.assertEqual(_run_count(renamed, text), 0)
|
|
assert_closed(self, renamed)
|
|
|
|
def test_a_write_racing_a_failed_switch_lands_in_the_old_library(self):
|
|
text = "a prompt racing a failed switch"
|
|
self.db.save_prompt(text, prompt_hash=generate_prompt_hash(text))
|
|
self.db.model.close_all()
|
|
writer_db = PromptDatabase(self.old_db)
|
|
thread, started, probe, outcome = self._racing_run_count(writer_db, text)
|
|
|
|
def checkpoint(step):
|
|
if step == "locked":
|
|
thread.start()
|
|
self.assertTrue(started.wait(WAIT))
|
|
|
|
def switch(old_db, location, detached):
|
|
lm.switch_process(old_db, location, detached)
|
|
raise RuntimeError("switch failed at its very end")
|
|
|
|
with probe.watching():
|
|
with self.assertRaises(lm.MoveFailed):
|
|
self.move(checkpoint=checkpoint, switch=switch)
|
|
# The copy was removed: no connection was left open on it.
|
|
self.assert_target_clean()
|
|
assert_closed(self, self.target)
|
|
thread.join(WAIT)
|
|
self.assertFalse(thread.is_alive())
|
|
self.assertEqual(probe.violations, [])
|
|
self.assertEqual(outcome, {"ok": True})
|
|
self.assertEqual(_run_count(self.old_db, text), 1)
|
|
|
|
def test_a_retire_failure_after_a_complete_switch_does_not_fail_the_move(self):
|
|
node = PromptDatabase(self.old_db)
|
|
node.get_prompt_by_id(1)
|
|
real = models.PromptModel.retire_detached
|
|
calls = []
|
|
|
|
def failing(connections, logger=None):
|
|
calls.append(len(connections))
|
|
real(connections, logger) # closed, but the call still fails
|
|
raise RuntimeError("retiring failed")
|
|
|
|
with mock.patch.object(
|
|
models.PromptModel, "retire_detached", staticmethod(failing)
|
|
):
|
|
with self.assertLogs("prompt_manager.database.library_move", "ERROR"):
|
|
result = self.move()
|
|
self.assertEqual(len(calls), 1)
|
|
self.assertGreater(calls[0], 0)
|
|
self.assertEqual(node.model.db_path, self.new_db)
|
|
self.assertEqual(result.location.db_path, self.new_db)
|
|
self.assertTrue(
|
|
os.path.isfile(
|
|
os.path.join(self.ext, "prompts.db.moved-to-__promptmanager")
|
|
)
|
|
)
|
|
|
|
@staticmethod
|
|
def _racing_run_count(writer_db, text):
|
|
"""A run count whose UPDATE marks its connection in use for the probe."""
|
|
started = threading.Event()
|
|
probe = _CrossThreadCloseProbe()
|
|
outcome = {}
|
|
|
|
def write():
|
|
conn = writer_db.model.get_connection()
|
|
|
|
def trace(sql):
|
|
if sql.startswith("UPDATE"):
|
|
probe.using(conn)
|
|
started.set()
|
|
|
|
conn.set_trace_callback(trace)
|
|
try:
|
|
usage_tracking.count_runs(
|
|
[(generate_prompt_hash(text), text)],
|
|
writer_db,
|
|
usage_tracking.PendingFirstUse(),
|
|
)
|
|
outcome["ok"] = True
|
|
except Exception as error: # pragma: no cover - reported below
|
|
outcome["error"] = error
|
|
finally:
|
|
writer_db.model.close()
|
|
|
|
return threading.Thread(target=write), started, probe, outcome
|
|
|
|
|
|
class TestVerifyUri(unittest.TestCase):
|
|
"""M2: a copy whose path holds ``#`` and ``%`` verifies."""
|
|
|
|
def test_copy_path_with_uri_characters_verifies(self):
|
|
import tempfile
|
|
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
folder = os.path.join(tmp, "lib #1 %20")
|
|
os.makedirs(folder)
|
|
path = os.path.join(folder, "copy.db")
|
|
db = PromptDatabase(path)
|
|
db.save_prompt("a prompt")
|
|
db.model.close_all()
|
|
with contextlib.closing(sqlite3.connect(path)) as conn:
|
|
conn.execute("PRAGMA journal_mode=DELETE")
|
|
counts, documents = lm._describe(conn)
|
|
lm._verify_copy(path, counts, documents)
|
|
models.PromptModel.close_all_instances()
|
|
|
|
|
|
class TestRestartFinishesTheRename(_MoveCase):
|
|
"""M6: the marker names the old files; a restart renames them."""
|
|
|
|
def _crash_after_switch(self):
|
|
def checkpoint(step):
|
|
if step == "switched":
|
|
raise SimulatedCrash(step)
|
|
|
|
with self.assertRaises(SimulatedCrash):
|
|
self.move(checkpoint=checkpoint)
|
|
models.PromptModel.close_all_instances()
|
|
ll.clear_override()
|
|
|
|
def test_a_restart_renames_the_old_files_never_deletes(self):
|
|
self._crash_after_switch()
|
|
self.assertTrue(os.path.isfile(self.old_db))
|
|
restarted = self.location()
|
|
self.assertEqual(restarted.db_path, self.new_db)
|
|
marker = ll.read_marker(self.target)
|
|
self.assertEqual(
|
|
(marker["from"], marker["config_from"]), (self.old_db, self.config)
|
|
)
|
|
renamed = lm.finish_interrupted_move(restarted)
|
|
self.assertEqual(
|
|
sorted(renamed),
|
|
[
|
|
"config.json.moved-to-__promptmanager",
|
|
"prompts.db.moved-to-__promptmanager",
|
|
],
|
|
)
|
|
self.assertFalse(os.path.exists(self.old_db))
|
|
moved_config = os.path.join(self.ext, "config.json.moved-to-__promptmanager")
|
|
with open(moved_config, encoding="utf-8") as handle:
|
|
self.assertEqual(
|
|
json.load(handle)["integrations"]["lora_manager"]["civitai_api_key"],
|
|
KEY,
|
|
)
|
|
self.assertEqual(lm.finish_interrupted_move(self.location()), [])
|
|
|
|
def test_nothing_is_renamed_while_the_old_file_is_in_use(self):
|
|
self._crash_after_switch()
|
|
reader = sqlite3.connect(self.old_db)
|
|
try:
|
|
reader.execute("BEGIN")
|
|
reader.execute("SELECT COUNT(*) FROM prompts").fetchone()
|
|
PromptDatabase(self.old_db).save_prompt("in the wal")
|
|
self.assertEqual(lm.finish_interrupted_move(self.location()), [])
|
|
self.assertTrue(os.path.isfile(self.old_db))
|
|
finally:
|
|
reader.rollback()
|
|
reader.close()
|
|
models.PromptModel.close_all_instances()
|
|
|
|
def test_a_node_folder_library_has_nothing_to_finish(self):
|
|
self.assertEqual(lm.finish_interrupted_move(self.location()), [])
|
|
self.assertTrue(os.path.isfile(self.old_db))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|