diff --git a/database/library_move.py b/database/library_move.py index 896a4bb..269a869 100644 --- a/database/library_move.py +++ b/database/library_move.py @@ -438,8 +438,8 @@ def switch_process(old_db: str, location, detached: list) -> None: """Point this process at ``location``: resolver, every model, observers. Adds the old file's connections, detached and still open, to - ``detached`` as soon as they are detached, so the mover closes them - (``PromptModel.close_detached``) whether or not the rest succeeds. + ``detached`` as soon as they are detached, so the mover retires them + (``PromptModel.retire_detached``) whether or not the rest succeeds. """ from .execution_history import rekey_history_restore from .review_changes import rekey_review_changes @@ -474,6 +474,8 @@ class _Move: self.published: List[str] = [] self.marker: Optional[str] = None self.moved = False + # Retired connections still in use; awaited once the locks are free. + self.retiring: list = [] def temp(self, suffix: str) -> str: path = os.path.join( @@ -704,7 +706,9 @@ def _switch_guarded(move: _Move, location) -> list: The copy's write lock keeps every repointed thread from writing to it until the switch is whole, so undoing a partial switch loses nothing. - Returns the old file's detached connections (still open). + Returns the old file's detached connections (still open). On failure + every detached connection is retired; those still in use are added to + ``move.retiring`` for the caller to await once the locks are released. """ plan = move.plan guard = sqlite3.connect( @@ -718,15 +722,19 @@ def _switch_guarded(move: _Move, location) -> list: except Exception as error: logger.exception("Library move: switching failed; undoing it") move.moved = False # the copy and its marker are removed + # Retire, never wait, here: a connection still in use may be + # waiting on the very locks this thread holds. try: - PromptModel.close_detached( - move.unswitch(plan.db_path, location.db_path), logger + move.retiring.extend( + PromptModel.retire_detached( + move.unswitch(plan.db_path, location.db_path), logger + ) ) except Exception: logger.exception("Library move: undoing the switch failed") # The old file's connections a partial switch detached: the # models reopen it on next use, so close them now, not at GC. - PromptModel.close_detached(detached, logger) + move.retiring.extend(PromptModel.retire_detached(detached, logger)) raise MoveFailed("switch-failed") from error move.checkpoint("switched") return detached @@ -737,6 +745,20 @@ def _switch_guarded(move: _Move, location) -> list: def _copy_and_switch(move: _Move, location, attempts: int): + """Copy, then switch under both locks; retired connections are awaited. + + A connection that was inside a statement when it was retired is closed + by its own thread once the statement ends, which may need the locks + released first; the old files are renamed (and a failed copy removed) + only after that, since Windows refuses to rename or delete an open file. + """ + try: + return _copy_and_switch_locked(move, location, attempts) + finally: + PromptModel.await_detached(move.retiring, logger) + + +def _copy_and_switch_locked(move: _Move, location, attempts: int): plan = move.plan for attempt in range(1, attempts + 1): source = sqlite3.connect( @@ -756,9 +778,10 @@ def _copy_and_switch(move: _Move, location, attempts: int): files = move.publish(temp_db, copies) detached = _switch_guarded(move, location) # Still holding the old file's lock: a write that was waiting - # on it fails when its connection closes (and runs again on - # the new library) instead of landing in the old file. - PromptModel.close_detached(detached, logger) + # on it is interrupted and cannot commit, so it fails (and + # runs again on the new library) instead of landing in the + # old file. Its own thread closes the connection. + move.retiring.extend(PromptModel.retire_detached(detached, logger)) finally: if source.in_transaction: source.execute("ROLLBACK") diff --git a/database/managed_connection.py b/database/managed_connection.py new file mode 100644 index 0000000..4eda698 --- /dev/null +++ b/database/managed_connection.py @@ -0,0 +1,260 @@ +"""A SQLite connection that is never closed under the thread using it. + +``PromptModel`` gives each thread its own connection, but Move library and +``close_all`` must retire every thread's connection from one thread. Closing +a ``sqlite3.Connection`` while its owner is inside a statement (Python runs +``sqlite3_step`` without the GIL) frees memory the owner then reads: on +Windows the process dies with "access violation", elsewhere the owner sees +garbage such as a ``MemoryError``. + +``ManagedConnection`` counts the calls its owner is inside. ``retire()`` is +safe from any thread: an idle connection is closed at once; one in use is +interrupted (``sqlite3_interrupt`` is thread-safe) and closed by its owner as +the call ends. A retired connection refuses new work with the error a closed +connection raises, and never commits: a write that started before the +retirement and waited for the database lock is rolled back, so it cannot +land in a file the process has switched away from. +""" + +from __future__ import annotations + +import contextlib +import logging +import threading +import time +import weakref +from typing import Iterable, List, Optional + +import sqlite3 + +CLOSED_MESSAGE = "Cannot operate on a closed database." + + +class ManagedCursor(sqlite3.Cursor): + """A cursor whose calls into SQLite count as its connection being in use.""" + + def execute(self, sql, parameters=(), /): + conn = self.connection + conn._enter() + try: + return super().execute(sql, parameters) + finally: + conn._leave() + + def executemany(self, sql, seq_of_parameters, /): + conn = self.connection + conn._enter() + try: + return super().executemany(sql, seq_of_parameters) + finally: + conn._leave() + + def executescript(self, sql_script, /): + conn = self.connection + conn._enter() + try: + return super().executescript(sql_script) + finally: + conn._leave() + + def fetchone(self): + conn = self.connection + conn._enter() + try: + return super().fetchone() + finally: + conn._leave() + + def fetchmany(self, *args, **kwargs): + conn = self.connection + conn._enter() + try: + return super().fetchmany(*args, **kwargs) + finally: + conn._leave() + + def fetchall(self): + conn = self.connection + conn._enter() + try: + return super().fetchall() + finally: + conn._leave() + + def __next__(self): + conn = self.connection + conn._enter() + try: + return super().__next__() + finally: + conn._leave() + + def close(self): + conn = self.connection + try: + conn._enter() + except sqlite3.ProgrammingError: + return super().close() # the connection is closed or closing + try: + return super().close() + finally: + conn._leave() + + +class ManagedConnection(sqlite3.Connection): + """A connection only its owner closes while a call into it is running.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._use_lock = threading.Lock() + self._depth = 0 + self._retired = False + self._released = threading.Event() + # Python before 3.12 keeps the file open after close() while a cursor + # still holds a statement; these are closed first (see _close_now). + self._cursors: "weakref.WeakSet[sqlite3.Cursor]" = weakref.WeakSet() + self.close_logger: Optional[logging.Logger] = None + + # ── in use ─────────────────────────────────────────────────────── + + def _enter(self) -> None: + with self._use_lock: + if self._retired: + raise sqlite3.ProgrammingError(CLOSED_MESSAGE) + self._depth += 1 + + def _leave(self) -> None: + with self._use_lock: + self._depth -= 1 + close_now = self._retired and self._depth == 0 + if close_now: + self._close_now() + + def cursor(self, factory=None): + cursor = super().cursor(factory or ManagedCursor) + self._cursors.add(cursor) + return cursor + + def execute(self, sql, parameters=(), /): + return self.cursor().execute(sql, parameters) + + def executemany(self, sql, seq_of_parameters, /): + return self.cursor().executemany(sql, seq_of_parameters) + + def executescript(self, sql_script, /): + return self.cursor().executescript(sql_script) + + def commit(self): + self._enter() + try: + return super().commit() + finally: + self._leave() + + def rollback(self): + self._enter() # a closed connection raises, as sqlite3's does + try: + return super().rollback() + finally: + self._leave() + + def backup(self, target, **kwargs): + self._enter() + try: + return super().backup(target, **kwargs) + finally: + self._leave() + + def __exit__(self, exc_type, exc_value, traceback): + # The C implementation calls its own commit/rollback, not these. + if exc_type is None: + try: + self.commit() + except BaseException: + with contextlib.suppress(sqlite3.Error): + self.rollback() + raise + else: + self.rollback() + return False + + # ── closing ────────────────────────────────────────────────────── + + def close(self) -> None: + """Close now if idle, else as soon as the running call ends (any thread).""" + self.retire() + + def retire(self) -> bool: + """Refuse new work and close; True when the connection is closed already.""" + with self._use_lock: + if not self._retired: + self._retired = True + if self._depth: + # Aborts the running statement, or the one waiting for + # the lock as soon as it gets it. Thread-safe; the owner + # cannot close the connection while this lock is held. + self.interrupt() + return False + elif self._depth: + return False + else: + return self._released.wait(0) + self._close_now() + return True + + def wait_closed(self, timeout: float) -> bool: + return self._released.wait(timeout) + + def _close_now(self) -> None: + # Only ever called with no call into the connection running, and a + # retired connection refuses new ones: its cursors are idle. + for cursor in list(self._cursors): + with contextlib.suppress(sqlite3.Error): + sqlite3.Cursor.close(cursor) + try: + super().close() + except sqlite3.Error as error: + if self.close_logger is not None: + self.close_logger.warning(f"Error closing database connection: {error}") + finally: + self._released.set() + + +def retire_connections( + connections: Iterable[sqlite3.Connection], logger=None +) -> List[ManagedConnection]: + """Retire every connection; return those whose owner has yet to close them. + + Safe while holding a lock the owners wait on: nothing here waits for them. + """ + pending = [] + for conn in connections: + if isinstance(conn, ManagedConnection): + if logger is not None: + conn.close_logger = logger + if not conn.retire(): + pending.append(conn) + continue + try: + conn.close() + except sqlite3.Error as error: + if logger is not None: + logger.warning(f"Error closing database connection: {error}") + return pending + + +def await_closed( + connections: Iterable[ManagedConnection], timeout: float, logger=None +) -> bool: + """Wait until the owners closed ``connections`` (retired); False on timeout.""" + deadline = time.monotonic() + timeout + still_open = 0 + for conn in connections: + if not conn.wait_closed(max(0.0, deadline - time.monotonic())): + still_open += 1 + if still_open and logger is not None: + logger.warning( + f"{still_open} database connection(s) still in use after " + f"{timeout:g}s; they close when their statement ends" + ) + return not still_open diff --git a/database/models.py b/database/models.py index 0c0210e..443d4ac 100644 --- a/database/models.py +++ b/database/models.py @@ -15,6 +15,11 @@ from contextlib import closing from pathlib import Path from typing import Optional, Set, Tuple +from .managed_connection import ( + ManagedConnection, + await_closed, + retire_connections, +) from .library_schema import library_schema_error, migrate_library_state from .image_favorites import image_favorites_schema_error, migrate_image_favorites from .annotation_schema import annotation_schema_error, migrate_annotation_state @@ -615,10 +620,15 @@ class PromptModel: # meanwhile: otherwise it would write to the old file. with self._conn_lock: path, generation = self.db_path, self._generation - # check_same_thread=False only so close_all can clean up other - # threads' connections; each connection is still used by one thread. + # check_same_thread=False only so close_all and Move library can + # retire other threads' connections; each connection is still + # used by one thread, and ManagedConnection never closes it under + # a running statement (see managed_connection). conn = sqlite3.connect( - path, check_same_thread=False, timeout=BUSY_TIMEOUT_SECONDS + path, + check_same_thread=False, + timeout=BUSY_TIMEOUT_SECONDS, + factory=ManagedConnection, ) try: conn.row_factory = sqlite3.Row @@ -654,9 +664,9 @@ class PromptModel: Every thread opens the new file on its next ``get_connection``. The old connections are detached, not closed, and returned: the caller - closes them once nothing can block on the new file any more. Closing - a connection another thread is using waits for its statement; a write - waiting on the mover's lock then fails (and, through + retires them (``retire_detached``) while it still holds the old + file's write lock. A write waiting on that lock is then interrupted + and can never commit, so it fails (and, through ``PromptDatabase.survive_library_switch``, runs again on the new file). """ old_key = os.path.normcase(os.path.realpath(old_path)) @@ -671,14 +681,31 @@ class PromptModel: return detached @staticmethod - def close_detached(connections, logger=None) -> None: - """Close connections ``repoint_instances`` detached.""" - for conn in connections: - try: - conn.close() - except sqlite3.Error as e: - if logger is not None: - logger.warning(f"Error closing database connection: {e}") + def retire_detached(connections, logger=None) -> list: + """Retire detached connections without waiting for their owners. + + Idle ones close now. One whose owner thread is inside a statement is + never closed from here (that frees memory the owner is using: an + access violation on Windows); it is interrupted, refuses to commit and + is closed by its owner as the statement ends. Returns those still + open, for ``await_detached`` once no lock the owners wait on is held. + """ + return retire_connections(connections, logger) + + @staticmethod + def await_detached(connections, logger=None) -> bool: + """Wait until retired connections are closed; False if some are not.""" + return await_closed(connections, BUSY_TIMEOUT_SECONDS, logger) + + @staticmethod + def close_detached(connections, logger=None) -> bool: + """Retire connections ``repoint_instances`` detached and await them. + + Do not call it while holding a lock their owners may wait on; retire + them there (``retire_detached``) and await them after releasing it. + """ + pending = retire_connections(connections, logger) + return await_closed(pending, BUSY_TIMEOUT_SECONDS, logger) def close_all(self) -> None: """Close every thread's connection; they reopen lazily on next use. @@ -714,10 +741,7 @@ class PromptModel: self._local.conn = None with self._conn_lock: self._connections.discard(conn) - try: - conn.close() - except sqlite3.Error as e: - self.logger.warning(f"Error closing database connection: {e}") + self.close_detached([conn], self.logger) def _migrate_workflow_name_removal(self, conn: sqlite3.Connection) -> None: """ diff --git a/tests/test_database_backup.py b/tests/test_database_backup.py index 7eddca7..e4aa9a8 100644 --- a/tests/test_database_backup.py +++ b/tests/test_database_backup.py @@ -293,6 +293,8 @@ class TestRestoreFromFile(BackupTestCase): super().backup(target, *args, **kwargs) def connect_with_failure(*args, **kwargs) -> sqlite3.Connection: + if "factory" in kwargs: # the model's own managed connections + return connect(*args, **kwargs) return connect(*args, **kwargs, factory=FailingLiveBackup) with patch("database.models.sqlite3.connect", side_effect=connect_with_failure): diff --git a/tests/test_library_move_review.py b/tests/test_library_move_review.py index 1e4fa82..4cfaa0f 100644 --- a/tests/test_library_move_review.py +++ b/tests/test_library_move_review.py @@ -7,6 +7,9 @@ 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. @@ -46,6 +49,39 @@ def _run_count(path, 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.""" @@ -166,44 +202,46 @@ class TestBackgroundWritesSurviveTheSwitch(_MoveCase): self.db.model.close_all() image = os.path.join(self.comfy, "output", "late.png") - # A short busy timeout so the blocked writer gives up in seconds. - with mock.patch.object(models, "BUSY_TIMEOUT_SECONDS", 2): - writer_db = PromptDatabase(self.old_db) - handler = ImageGenerationHandler(writer_db, mock.MagicMock()) - started = threading.Event() - outcome = {} + writer_db = PromptDatabase(self.old_db) + handler = ImageGenerationHandler(writer_db, mock.MagicMock()) + started = threading.Event() + probe = _CrossThreadCloseProbe() + outcome = {} - def write(): - conn = writer_db.model.get_connection() - conn.set_trace_callback( - lambda sql: ( - started.set() - if sql.lstrip().upper().startswith(("BEGIN", "INSERT")) - else None - ) - ) - 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() + def write(): + conn = writer_db.model.get_connection() - thread = threading.Thread(target=write) + def trace(sql): + if sql.lstrip().upper().startswith(("BEGIN", "INSERT")): + probe.using(conn) + started.set() - 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)) + 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() - self.move(checkpoint=checkpoint) - thread.join(WAIT) + 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"] @@ -229,45 +267,87 @@ class TestBackgroundWritesSurviveTheSwitch(_MoveCase): text = "a racing prompt" self.db.save_prompt(text, prompt_hash=generate_prompt_hash(text)) self.db.model.close_all() - # A short busy timeout so the blocked writer gives up in seconds. - with mock.patch.object(models, "BUSY_TIMEOUT_SECONDS", 2): - writer_db = PromptDatabase(self.old_db) - started = threading.Event() - outcome = {} + writer_db = PromptDatabase(self.old_db) + thread, started, probe, outcome = self._racing_run_count(writer_db, text) - def write(): - conn = writer_db.model.get_connection() - conn.set_trace_callback( - lambda sql: started.set() if sql.startswith("UPDATE") else None - ) - 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() + 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)) - 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 UPDATE is running (it then waits on us). - thread.start() - self.assertTrue(started.wait(WAIT)) - - self.move(checkpoint=checkpoint) - thread.join(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) + + @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): diff --git a/tests/test_managed_connection.py b/tests/test_managed_connection.py new file mode 100644 index 0000000..8303623 --- /dev/null +++ b/tests/test_managed_connection.py @@ -0,0 +1,220 @@ +"""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 +from unittest import 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_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 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()