Files
ComfyAssets-ComfyUI_PromptM…/tests/test_managed_connection.py
T
Vito d88cf9e7c6 fix(quality): sweep the Code Quality standard findings (#558)
* 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
2026-10-31 05:43:48 +00:00

443 lines
17 KiB
Python

"""ManagedConnection: never closed under the thread using it.
Closing a sqlite3 connection from one thread while another is inside a
statement on it is a use-after-free in SQLite ("Windows fatal exception:
access violation" in CI). A retired connection is closed at once when idle,
otherwise interrupted and closed by its owner as the call ends; it never
commits again.
"""
import contextlib
import os
import sqlite3
import sys
import tempfile
import threading
import unittest
import weakref
import unittest.mock as mock
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from database.managed_connection import ( # noqa: E402
ManagedConnection,
ManagedCursor,
await_closed,
retire_connections,
)
from database.models import PromptModel # noqa: E402
from tests.open_handles import assert_closed # noqa: E402
WAIT = 30
class _Case(unittest.TestCase):
def setUp(self):
self._tmp = tempfile.TemporaryDirectory()
self.path = os.path.join(self._tmp.name, "managed.db")
with contextlib.closing(sqlite3.connect(self.path)) as conn:
conn.execute("PRAGMA journal_mode=WAL")
conn.execute("CREATE TABLE t(x INTEGER)")
conn.execute("INSERT INTO t VALUES (0)")
conn.commit()
self.opened = []
def tearDown(self):
for conn in self.opened:
conn.close()
try:
assert_closed(self, self._tmp.name)
finally:
self._tmp.cleanup()
def connect(self, **kwargs):
conn = sqlite3.connect(
self.path,
check_same_thread=False,
timeout=WAIT,
factory=ManagedConnection,
**kwargs,
)
self.opened.append(conn)
return conn
def value(self):
with contextlib.closing(sqlite3.connect(self.path)) as conn:
return conn.execute("SELECT x FROM t").fetchone()[0]
class TestRetire(_Case):
def test_an_idle_connection_retired_from_another_thread_closes_at_once(self):
conn = self.connect()
conn.execute("SELECT 1").fetchone()
closed = []
thread = threading.Thread(target=lambda: closed.append(conn.retire()))
thread.start()
thread.join(WAIT)
self.assertEqual(closed, [True])
self.assertTrue(conn.wait_closed(0))
with self.assertRaises(sqlite3.ProgrammingError):
conn.execute("SELECT 1")
self.assertTrue(conn.retire()) # retiring again is harmless
def test_a_connection_in_use_is_closed_by_its_owner_when_the_call_ends(self):
conn = self.connect()
holder = sqlite3.connect(self.path, isolation_level=None)
self.opened.append(holder)
holder.execute("BEGIN IMMEDIATE")
started = threading.Event()
outcome = {}
conn.set_trace_callback(
lambda sql: started.set() if sql.startswith("UPDATE") else None
)
def write():
try:
with conn:
conn.execute("UPDATE t SET x = 1")
outcome["ok"] = True
except sqlite3.Error as error:
outcome["error"] = error
thread = threading.Thread(target=write)
thread.start()
self.assertTrue(started.wait(WAIT))
# The owner is waiting on our lock inside its UPDATE: not closed here.
self.assertFalse(conn.retire())
self.assertFalse(conn.retire())
self.assertFalse(conn.wait_closed(0))
self.assertEqual(retire_connections([conn]), [conn])
holder.execute("ROLLBACK")
thread.join(WAIT)
self.assertFalse(thread.is_alive())
self.assertTrue(conn.wait_closed(WAIT))
self.assertIsInstance(outcome.get("error"), sqlite3.Error)
self.assertEqual(self.value(), 0)
def test_a_write_retired_before_its_commit_is_rolled_back(self):
conn = self.connect()
conn.execute("UPDATE t SET x = 2") # an implicit transaction is open
thread = threading.Thread(target=conn.retire)
thread.start()
thread.join(WAIT)
with self.assertRaises(sqlite3.ProgrammingError):
conn.commit()
with self.assertRaises(sqlite3.ProgrammingError):
with conn:
pass
with self.assertRaises(sqlite3.ProgrammingError):
conn.rollback()
self.assertEqual(self.value(), 0)
def test_close_is_retire(self):
conn = self.connect()
conn.close()
self.assertTrue(conn.wait_closed(0))
with self.assertRaises(sqlite3.ProgrammingError):
conn.commit()
def test_plain_connections_are_closed_directly(self):
plain = sqlite3.connect(self.path)
self.assertEqual(retire_connections([plain]), [])
with self.assertRaises(sqlite3.ProgrammingError):
plain.execute("SELECT 1")
def test_awaiting_a_connection_still_in_use_times_out_with_a_warning(self):
conn = self.connect()
conn._enter() # its owner is inside a call
self.assertEqual(retire_connections([conn]), [conn])
logger = mock.Mock()
self.assertFalse(await_closed([conn], 0, logger))
logger.warning.assert_called_once()
conn._leave()
self.assertTrue(await_closed([conn], WAIT, logger))
class TestCalls(_Case):
def test_iteration_and_fetches_share_one_position(self) -> None:
conn = self.connect()
conn.executemany("INSERT INTO t VALUES (?)", [(x,) for x in range(1, 260)])
cursor = conn.execute("SELECT x FROM t ORDER BY x")
self.assertIs(iter(cursor), cursor)
self.assertEqual(next(cursor), (0,))
self.assertEqual(cursor.fetchone(), (1,))
self.assertEqual(cursor.fetchmany(2), [(2,), (3,)])
self.assertEqual(next(iter(cursor)), (4,))
self.assertEqual(cursor.fetchall(), [(x,) for x in range(5, 260)])
self.assertEqual(list(cursor), [])
def test_a_partial_loop_can_resume_without_skipping_rows(self) -> None:
cursor = self.connect().execute(
"SELECT 1 UNION ALL SELECT 2 UNION ALL SELECT 3"
)
for row in cursor:
self.assertEqual(row, (1,))
break
self.assertEqual(list(cursor), [(2,), (3,)])
def test_an_existing_iterator_refuses_reads_after_retire_or_close(self) -> None:
for retire in (False, True):
with self.subTest(retire=retire):
conn = self.connect()
cursor = conn.execute("SELECT 1 UNION ALL SELECT 2")
iterator = iter(cursor)
self.assertEqual(next(iterator), (1,))
if retire:
conn.retire()
else:
cursor.close()
with self.assertRaises(sqlite3.ProgrammingError):
next(iterator)
def test_buffered_fetchmany_matches_sqlite_sizes_and_validation(self) -> None:
conn = self.connect()
conn.executemany("INSERT INTO t VALUES (?)", [(x,) for x in range(1, 260)])
conn.commit()
with contextlib.closing(sqlite3.connect(self.path)) as plain:
for size in (None, 0, 1, 2, 127, 128, 130, 300, -1):
with self.subTest(size=size):
cursors = [
c.execute("SELECT x FROM t ORDER BY x") for c in (plain, conn)
]
for cursor in cursors:
cursor.arraysize = 2
self.assertEqual(next(iter(cursor)), (0,))
results = []
for cursor in cursors:
try:
rows = (
cursor.fetchmany()
if size is None
else cursor.fetchmany(size)
)
except ValueError as error: # Negative sizes on Python 3.14+.
rows = type(error)
results.append((rows, cursor.fetchall()))
self.assertEqual(*results)
for args, kwargs in ((("bad",), {}), ((1, 2), {}), ((), {"unknown": 1})):
cursor = conn.execute("SELECT 1 UNION ALL SELECT 2")
self.assertEqual(next(iter(cursor)), (1,))
with self.assertRaises(TypeError):
cursor.fetchmany(*args, **kwargs)
self.assertEqual(cursor.fetchone(), (2,))
def test_reexecuting_a_cursor_matches_sqlite_result_lifetime(self) -> None:
conn = self.connect()
conn.row_factory = sqlite3.Row
for method, args in (
("execute", ("SELECT 9 AS x",)),
("executemany", ("INSERT INTO t VALUES (?)", [(9,)])),
("executescript", ("SELECT 9 AS x;",)),
):
with self.subTest(method=method):
cursor = conn.execute("SELECT 1 AS x UNION ALL SELECT 2")
self.assertEqual(next(iter(cursor))["x"], 1)
self.assertIs(getattr(cursor, method)(*args), cursor)
if method == "execute":
self.assertEqual(cursor.fetchone()["x"], 9)
elif method == "executescript":
self.assertEqual(cursor.fetchone()["x"], 2)
self.assertEqual(list(cursor), [])
def test_failed_execution_matches_sqlite_result_lifetime(self) -> None:
conn = self.connect()
cases = (
("execute", (None,)),
("execute", ("invalid sql",)),
("execute", ("SELECT ?", ())),
("executemany", (None, [])),
("executemany", ("SELECT 9", [])),
("executescript", (None,)),
("executescript", ("invalid sql;",)),
)
with contextlib.closing(sqlite3.connect(self.path)) as plain:
for method, args in cases:
with self.subTest(method=method, args=args):
results = []
for connection in (plain, conn):
cursor = connection.execute("SELECT 1 UNION ALL SELECT 2")
self.assertEqual(next(iter(cursor)), (1,))
with self.assertRaises(
sqlite3.Error if isinstance(args[0], str) else TypeError
):
getattr(cursor, method)(*args)
results.append(cursor.fetchall())
self.assertEqual(*results)
def test_every_cursor_call_works_and_is_refused_once_retired(self):
conn = self.connect()
conn.row_factory = sqlite3.Row
cursor = conn.cursor()
self.assertIsInstance(cursor, ManagedCursor)
conn.executemany("INSERT INTO t VALUES (?)", [(1,), (2,)])
conn.executescript("INSERT INTO t VALUES (3);")
self.assertEqual(
[row["x"] for row in conn.execute("SELECT x FROM t ORDER BY x")],
[0, 1, 2, 3],
)
cursor.execute("SELECT x FROM t ORDER BY x")
self.assertEqual(cursor.fetchone()["x"], 0)
self.assertEqual([r["x"] for r in cursor.fetchmany(2)], [1, 2])
self.assertEqual([r["x"] for r in cursor.fetchall()], [3])
cursor.executemany("UPDATE t SET x = x WHERE x = ?", [(1,)])
cursor.executescript("SELECT 1;")
cursor.close()
conn.commit()
conn.rollback()
with conn:
conn.execute("DELETE FROM t WHERE x = 3")
with self.assertRaises(ZeroDivisionError):
with conn:
conn.execute("DELETE FROM t")
raise ZeroDivisionError
self.assertEqual(self.value(), 0)
copy = sqlite3.connect(":memory:")
self.opened.append(copy)
conn.backup(copy)
self.assertEqual(copy.execute("SELECT COUNT(*) FROM t").fetchone()[0], 3)
reader = conn.execute("SELECT x FROM t")
conn.retire()
for call in (
lambda: reader.fetchone(),
lambda: reader.fetchmany(1),
lambda: reader.fetchall(),
lambda: next(reader),
lambda: reader.execute("SELECT 1"),
lambda: reader.executemany("SELECT ?", [(1,)]),
lambda: reader.executescript("SELECT 1;"),
lambda: conn.backup(copy),
):
with self.assertRaises(sqlite3.ProgrammingError):
call()
# As with sqlite3's own cursors (the result varies by Python version).
with contextlib.suppress(sqlite3.ProgrammingError):
reader.close()
class _CursorOpenedDuringClose(weakref.WeakSet):
"""The cursor set of a connection being closed; a cursor opens mid-iteration.
``hook`` runs after the first cursor is yielded, where a thread switch
could let the owner call ``conn.cursor()`` (Python 3.10-3.12 copy a
WeakSet one item at a time). It runs in the closing thread itself, so
the result does not depend on scheduling.
"""
def __init__(self, items, hook):
super().__init__(items)
self.hook = hook
def __iter__(self):
hooked = False
for item in super().__iter__():
yield item
if not hooked:
hooked = True
self.hook()
class TestReviewRound(_Case):
"""PR #321 review: I1 (cursor set race), M1 and M2 (unguarded paths)."""
def test_a_cursor_opened_while_the_connection_closes_cannot_break_retire(self):
conn = self.connect()
kept = [conn.cursor() for _ in range(3)]
refused = []
def open_cursor():
try:
kept.append(conn.cursor())
except sqlite3.ProgrammingError as error:
refused.append(error)
conn._cursors = _CursorOpenedDuringClose(conn._cursors, open_cursor)
self.assertTrue(conn.retire()) # no "Set changed size during iteration"
self.assertTrue(conn.wait_closed(0))
self.assertEqual(len(refused), 1)
self.assertEqual(len(kept), 3)
def test_a_retired_connection_refuses_new_cursors_and_scripts(self):
conn = self.connect()
conn._enter() # its owner is inside a call: retired, not yet closed
try:
self.assertFalse(conn.retire())
for call in (
conn.cursor,
lambda: conn.executescript("CREATE TABLE late(x);"),
lambda: conn.in_transaction,
lambda: conn.create_function("f", 0, lambda: 1),
):
with self.subTest(call=call):
with self.assertRaises(sqlite3.ProgrammingError):
call()
finally:
conn._leave()
self.assertTrue(conn.wait_closed(0))
with contextlib.closing(sqlite3.connect(self.path)) as check:
tables = check.execute(
"SELECT name FROM sqlite_master WHERE name='late'"
).fetchall()
self.assertEqual(tables, [])
def test_one_connection_failing_to_retire_does_not_skip_the_rest(self):
class Broken(ManagedConnection):
def retire(self):
raise RuntimeError("retire failed")
broken = sqlite3.connect(self.path, check_same_thread=False, factory=Broken)
try:
after = self.connect()
logger = mock.Mock()
pending = retire_connections([broken, after], logger)
self.assertEqual(pending, [])
self.assertTrue(after.wait_closed(0))
logger.exception.assert_called_once()
finally:
ManagedConnection.retire(broken)
def test_a_restore_marks_its_backup_target_in_use(self):
model = PromptModel(os.path.join(self._tmp.name, "live.db"))
try:
self._restore_marks_its_target(model)
finally:
model.close_all()
def _restore_marks_its_target(self, model):
source_path = os.path.join(self._tmp.name, "source.db")
PromptModel(source_path).close_all()
target = model.get_connection()
depths = []
connect = sqlite3.connect
class Source(sqlite3.Connection):
def backup(self, into, *args, **kwargs):
depths.append(into._depth)
return super().backup(into, *args, **kwargs)
def connect_source(*args, **kwargs):
if "factory" in kwargs:
return connect(*args, **kwargs)
return connect(*args, **kwargs, factory=Source)
with mock.patch("database.models.sqlite3.connect", connect_source):
model._replace_contents_from(source_path)
self.assertEqual(depths, [1])
self.assertEqual(target._depth, 0)
class TestModelConnections(_Case):
def test_the_model_hands_out_managed_connections_and_closes_them(self):
model = PromptModel(os.path.join(self._tmp.name, "model.db"))
conn = model.get_connection()
self.assertIsInstance(conn, ManagedConnection)
model.close()
self.assertTrue(conn.wait_closed(0))
conn = model.get_connection()
model.close_all()
self.assertTrue(conn.wait_closed(0))
if __name__ == "__main__":
unittest.main()