From 856609a43a461d27f0288dc5b1f70a7b5df8c467 Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Fri, 30 Oct 2026 20:53:30 -0700 Subject: [PATCH] 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. --- database/managed_connection.py | 84 ++++++++++++++++------- py/api/autotag_routes.py | 4 +- tests/test_api_autotag.py | 27 ++++++++ tests/test_managed_connection.py | 110 +++++++++++++++++++++++++++++++ web/js/gallery.js | 1 - 5 files changed, 201 insertions(+), 25 deletions(-) diff --git a/database/managed_connection.py b/database/managed_connection.py index fb71fc6d..49e58390 100644 --- a/database/managed_connection.py +++ b/database/managed_connection.py @@ -35,38 +35,51 @@ from __future__ import annotations import contextlib import logging +import operator +import sys import threading import time import weakref -from typing import Iterable, Iterator, List, Optional +from collections import deque +from typing import Any, Callable, Iterable, Iterator, List, Optional import sqlite3 CLOSED_MESSAGE = "Cannot operate on a closed database." # Rows fetched per step when a cursor is iterated: the connection is marked # in use once per chunk instead of once per row (iteration costs about what -# a plain cursor costs). A loop that stops early leaves up to this many rows -# fetched; ``fetchone``/``fetchmany`` after a partial loop skip them. +# a plain cursor costs). Prefetched rows belong to the cursor so iteration +# and explicit fetches always advance the same position. ITER_CHUNK = 128 class ManagedCursor(sqlite3.Cursor): """A cursor whose calls into SQLite count as its connection being in use.""" + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self._pending = deque() + self._closed = False + def execute(self, sql, parameters=(), /): - conn = self.connection - conn._enter() - try: - return super().execute(sql, parameters) - finally: - conn._leave() + return self._run_statement(super().execute, sql, parameters) def executemany(self, sql, seq_of_parameters, /): + return self._run_statement(super().executemany, sql, seq_of_parameters) + + def _run_statement(self, operation: Callable, *args: Any) -> ManagedCursor: conn = self.connection conn._enter() + description = self.description + completed = False try: - return super().executemany(sql, seq_of_parameters) + result = operation(*args) + completed = True + return result finally: + # Argument parsing can fail before SQLite replaces the old statement. + if completed or self.description is not description: + self._pending.clear() conn._leave() def executescript(self, sql_script, /): @@ -81,6 +94,8 @@ class ManagedCursor(sqlite3.Cursor): conn = self.connection conn._enter() try: + if self._pending and not self._closed: + return self._pending.popleft() return super().fetchone() finally: conn._leave() @@ -89,7 +104,20 @@ class ManagedCursor(sqlite3.Cursor): conn = self.connection conn._enter() try: - return super().fetchmany(*args, **kwargs) + # Let SQLite validate the original arguments, including version-specific + # handling of zero/negative sizes, before touching prefetched rows. + rows = super().fetchmany(*args, **kwargs) + if not self._pending: + return rows + size = operator.index( + kwargs.get("size", args[0] if args else self.arraysize) + ) + # Python 3.12 reads all rows for size=0; 3.13+ reads none. + if size == 0 and sys.version_info >= (3, 13): + return rows + self._pending.extend(rows) + count = len(self._pending) if size <= 0 else min(size, len(self._pending)) + return [self._pending.popleft() for _ in range(count)] finally: conn._leave() @@ -97,25 +125,31 @@ class ManagedCursor(sqlite3.Cursor): conn = self.connection conn._enter() try: - return super().fetchall() + rows = super().fetchall() + pending = list(self._pending) + self._pending.clear() + return pending + rows finally: conn._leave() def __iter__(self) -> Iterator: - return self._chunks() - - def _chunks(self) -> Iterator: - while True: - rows = self.fetchmany(ITER_CHUNK) - if not rows: - return - yield from rows + return self def __next__(self): conn = self.connection + pending = self._pending + # Buffered Python rows need no SQLite lock, but retirement and cursor + # close must still stop an iterator obtained before either operation. + if conn._retired or self._closed: + raise sqlite3.ProgrammingError(CLOSED_MESSAGE) + if pending: + return pending.popleft() conn._enter() try: - return super().__next__() + pending.extend(super().fetchmany(ITER_CHUNK)) + if not pending: + raise StopIteration + return pending.popleft() finally: conn._leave() @@ -126,7 +160,10 @@ class ManagedCursor(sqlite3.Cursor): except sqlite3.ProgrammingError: return super().close() # the connection is closed or closing try: - return super().close() + result = super().close() + self._pending = deque() + self._closed = True + return result finally: conn._leave() @@ -293,6 +330,9 @@ class ManagedConnection(sqlite3.Connection): cursors = list(self._cursors) try: for cursor in cursors: + if isinstance(cursor, ManagedCursor): + cursor._pending = deque() + cursor._closed = True with contextlib.suppress(sqlite3.Error): sqlite3.Cursor.close(cursor) super().close() diff --git a/py/api/autotag_routes.py b/py/api/autotag_routes.py index 388e5e3a..d4038448 100644 --- a/py/api/autotag_routes.py +++ b/py/api/autotag_routes.py @@ -1041,12 +1041,12 @@ class AutotagRoutesMixin: return self._bad_request('existing must be "skip" or "merge"') try: options = self._single_options(data, model_type) + except OptionError as error: + return self._bad_request(str(error)) except (ValueError, TypeError): return self._bad_request( "general_threshold and character_threshold must be numeric" ) - except OptionError as error: - return self._bad_request(str(error)) target, error = await self._select_autotag_target(data) if error is not None: return error diff --git a/tests/test_api_autotag.py b/tests/test_api_autotag.py index 4b009ba8..886482bd 100644 --- a/tests/test_api_autotag.py +++ b/tests/test_api_autotag.py @@ -600,6 +600,33 @@ class TestAutotagSingleEdgeCases(AutotagAPITestCase): {"path": str(image), "general_threshold": "high"}, ) self.assertEqual(resp.status, 400) + self.assertEqual( + (await resp.json())["error"], + "general_threshold and character_threshold must be numeric", + ) + + async def test_invalid_single_options_report_the_validation_error(self) -> None: + image = self._make_image(self.output_dir) + cases = ( + ( + {"model_type": "wd14-vit", "general_threshold": 2}, + "general_threshold must be between 0 and 1", + ), + ( + {"model_type": "gguf", "prompt": "x" * 4001}, + "instructions must be at most 4000 characters", + ), + ) + for options, message in cases: + with self.subTest(options=options["model_type"]): + resp = await self._post_json( + "/prompt_manager/autotag/single", + {"path": str(image), **options}, + ) + self.assertEqual(resp.status, 400) + self.assertEqual( + await resp.json(), {"success": False, "error": message} + ) async def test_prompt_lookup_failure_still_returns_tags(self): image = self._make_image(self.output_dir) diff --git a/tests/test_managed_connection.py b/tests/test_managed_connection.py index b33d437a..7e1fbb38 100644 --- a/tests/test_managed_connection.py +++ b/tests/test_managed_connection.py @@ -154,6 +154,116 @@ class TestRetire(_Case): 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 diff --git a/web/js/gallery.js b/web/js/gallery.js index acb4b9a5..b889023b 100644 --- a/web/js/gallery.js +++ b/web/js/gallery.js @@ -607,7 +607,6 @@ button: true, backdrop: true, navbar: true, - title: true, // IMPORTANT: Use original images for viewing, not thumbnails url: 'data-original', // Tell ViewerJS to use data-original attribute toolbar: {