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: {