diff --git a/py/services/download_queue_service.py b/py/services/download_queue_service.py index 2f2fed89..a1ec500f 100644 --- a/py/services/download_queue_service.py +++ b/py/services/download_queue_service.py @@ -8,6 +8,7 @@ import sqlite3 import time from typing import Any, List, Optional +from ..utils.cache_db import closing_connection from ..utils.cache_paths import get_cache_base_dir logger = logging.getLogger(__name__) @@ -127,7 +128,7 @@ class DownloadQueueService: def _initialize_schema(self) -> None: if self._schema_initialized: return - with self._connect() as conn: + with closing_connection(self._connect()) as conn: conn.executescript(self._SCHEMA_TABLES) # Databases created by older versions lack diff --git a/py/services/downloaded_version_history_service.py b/py/services/downloaded_version_history_service.py index 7570af05..2c44ab4e 100644 --- a/py/services/downloaded_version_history_service.py +++ b/py/services/downloaded_version_history_service.py @@ -11,6 +11,7 @@ import sqlite3 import time from typing import Iterable, Mapping, Optional, Sequence +from ..utils.cache_db import closing_connection from ..utils.cache_paths import get_cache_base_dir from .settings_manager import get_settings_manager @@ -100,7 +101,7 @@ class DownloadedVersionHistoryService: def _initialize_schema(self) -> None: if self._schema_initialized: return - with self._connect() as conn: + with closing_connection(self._connect()) as conn: conn.executescript(self._SCHEMA) conn.commit() self._schema_initialized = True diff --git a/py/services/model_update_service.py b/py/services/model_update_service.py index d90ed9ab..3a80ba68 100644 --- a/py/services/model_update_service.py +++ b/py/services/model_update_service.py @@ -17,6 +17,7 @@ from typing import Any, Dict, Iterable, Iterator, List, Mapping, Optional, Seque from .errors import RateLimitError, ResourceNotFoundError from .settings_manager import get_settings_manager +from ..utils.cache_db import closing_connection from ..utils.cache_paths import CacheType, resolve_cache_path_with_migration from ..utils.constants import MODEL_WEIGHT_FILE_TYPES from ..utils.civitai_utils import rewrite_preview_url @@ -445,7 +446,7 @@ class ModelUpdateService: if self._schema_initialized: return try: - with self._connect() as conn: + with closing_connection(self._connect()) as conn: conn.execute("PRAGMA journal_mode=WAL") conn.execute("PRAGMA foreign_keys = ON") conn.executescript(self._SCHEMA) @@ -484,7 +485,7 @@ class ModelUpdateService: return try: - with sqlite3.connect(legacy_path, check_same_thread=False) as legacy_conn: + with closing_connection(sqlite3.connect(legacy_path, check_same_thread=False)) as legacy_conn: legacy_conn.row_factory = sqlite3.Row status_rows = legacy_conn.execute( """ @@ -2315,7 +2316,7 @@ class ModelUpdateService: status_rows: list[sqlite3.Row] = [] version_rows: list[sqlite3.Row] = [] - with self._connect() as conn: + with closing_connection(self._connect()) as conn: for start in range(0, len(ids), self._SQLITE_MAX_VARIABLES): chunk = tuple(ids[start : start + self._SQLITE_MAX_VARIABLES]) placeholders = ",".join("?" for _ in chunk) @@ -2418,7 +2419,7 @@ class ModelUpdateService: record.last_checked_at, 1 if record.should_ignore_model else 0, ) - with self._connect() as conn: + with closing_connection(self._connect()) as conn: conn.execute( """ INSERT INTO model_update_status ( diff --git a/py/services/persistent_model_cache.py b/py/services/persistent_model_cache.py index e7362f29..8008f07a 100644 --- a/py/services/persistent_model_cache.py +++ b/py/services/persistent_model_cache.py @@ -6,7 +6,7 @@ import threading from dataclasses import dataclass, field from typing import Any, Dict, List, Mapping, Optional, Sequence, Set, Tuple -from ..utils.cache_db import connect_cache_db +from ..utils.cache_db import closing_connection, connect_cache_db from ..utils.cache_paths import CacheType, resolve_cache_path_with_migration from ..utils.file_lock import exclusive_lock from ..utils.root_paths import normalize_root_identity @@ -618,7 +618,7 @@ class PersistentModelCache: if self._schema_initialized: return try: - with self._connect() as conn: + with closing_connection(self._connect()) as conn: conn.execute("PRAGMA journal_mode=WAL") conn.execute("PRAGMA foreign_keys = ON") conn.executescript( diff --git a/py/services/persistent_recipe_cache.py b/py/services/persistent_recipe_cache.py index 8b023135..e3bdbeda 100644 --- a/py/services/persistent_recipe_cache.py +++ b/py/services/persistent_recipe_cache.py @@ -19,7 +19,7 @@ import threading from dataclasses import dataclass, field from typing import Any, Dict, List, Optional, Set, Tuple -from ..utils.cache_db import connect_cache_db +from ..utils.cache_db import closing_connection, connect_cache_db from ..utils.cache_paths import CacheType, resolve_cache_path_with_migration from ..utils.file_lock import exclusive_lock @@ -457,7 +457,7 @@ class PersistentRecipeCache: if self._schema_initialized: return try: - with self._connect() as conn: + with closing_connection(self._connect()) as conn: conn.execute("PRAGMA journal_mode=WAL") conn.execute("PRAGMA foreign_keys = ON") conn.executescript( diff --git a/py/utils/cache_db.py b/py/utils/cache_db.py index 7afdb6a8..ffe1517a 100644 --- a/py/utils/cache_db.py +++ b/py/utils/cache_db.py @@ -14,7 +14,8 @@ the write path in WAL so readers are never blocked by a writer. from __future__ import annotations import sqlite3 -from typing import Any +from contextlib import contextmanager +from typing import Any, Iterator # How long a connection waits for a competing writer before raising. CONCURRENT_TIMEOUT_SECONDS = 30.0 @@ -79,3 +80,19 @@ def connect_cache_db( pass return conn + + +@contextmanager +def closing_connection(conn: sqlite3.Connection) -> Iterator[sqlite3.Connection]: + """Yield *conn* and close it on exit. + + ``with sqlite3.Connection`` only ends the transaction: it leaves the file + handle open, and a live handle keeps the cache database undeletable on + Windows (``WinError 32``) until the garbage collector happens to run. Every + ``with`` over a cache connection goes through here instead. + """ + + try: + yield conn + finally: + conn.close() diff --git a/tests/test_persistent_recipe_cache.py b/tests/test_persistent_recipe_cache.py index 5243da76..c01ba390 100644 --- a/tests/test_persistent_recipe_cache.py +++ b/tests/test_persistent_recipe_cache.py @@ -724,3 +724,33 @@ class TestHasWorkflowColumn: assert loaded is not None assert loaded.raw_data[0]["id"] == "legacy-1" assert loaded.raw_data[0]["has_workflow"] is False + + +def test_schema_initialization_closes_its_connection(tmp_path, monkeypatch): + """The schema init must not leak a file handle. + + ``with sqlite3.Connection`` only ends the transaction: the handle stays open, + and on Windows a live handle keeps the cache database undeletable + (``WinError 32``) until the garbage collector happens to run -- which broke + every teardown that removes a temporary cache file. + """ + + import sqlite3 + + from py.services import persistent_recipe_cache as module + + opened = [] + real_connect = sqlite3.connect + + def tracking_connect(*args, **kwargs): + conn = real_connect(*args, **kwargs) + opened.append(conn) + return conn + + monkeypatch.setattr(sqlite3, "connect", tracking_connect) + module.PersistentRecipeCache(db_path=str(tmp_path / "recipe.sqlite")) + + assert opened, "initializing the schema should have opened a connection" + for conn in opened: + with pytest.raises(sqlite3.ProgrammingError): + conn.execute("SELECT 1")