From a549690c7b7d693675aca71c59b6911ff2440a25 Mon Sep 17 00:00:00 2001 From: Copilot <223556219+Copilot@users.noreply.github.com> Date: Fri, 23 Oct 2026 14:37:02 -0400 Subject: [PATCH] FIX: Wait up to 30 seconds for file-backed SQLite locks File-backed SQLiteMemory engines inherited sqlite3's 5 second busy timeout, so concurrent scenario writes could fail with "database is locked" while another connection held SQLite's write lock. Configure a 30 second timeout for file-backed sync and async engines; in-memory engines are unchanged. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: 38248c3f-bada-4a90-beb2-be66c2ede072 --- doc/code/memory/0_memory.md | 4 +- pyrit/memory/sqlite_memory.py | 15 +++- .../memory/test_conversation_branching.py | 17 +--- tests/unit/memory/test_sqlite_memory.py | 88 ++++++++++++++++++- 4 files changed, 104 insertions(+), 20 deletions(-) diff --git a/doc/code/memory/0_memory.md b/doc/code/memory/0_memory.md index c5a533b0c5..c542bf9500 100644 --- a/doc/code/memory/0_memory.md +++ b/doc/code/memory/0_memory.md @@ -18,7 +18,9 @@ messages = await memory.get_conversation_messages_async(conversation_id="example Use the `_async` memory methods in async code. The synchronous methods remain available during deprecation. In-memory SQLite serializes transactions across -threads and event loops to prevent shared-cache table locks. +threads and event loops to prevent shared-cache table locks. File-backed SQLite +waits up to 30 seconds for another connection's lock before raising +`database is locked`. Repeated initialization reuses the existing memory without running schema setup again. Setup raises an error if `CentralMemory` and the requested backend singleton diff --git a/pyrit/memory/sqlite_memory.py b/pyrit/memory/sqlite_memory.py index 047c577f84..3fe096caa2 100644 --- a/pyrit/memory/sqlite_memory.py +++ b/pyrit/memory/sqlite_memory.py @@ -14,7 +14,7 @@ from sqlite3 import Connection as SQLiteConnection from sqlite3 import Cursor as SQLiteCursor from types import TracebackType -from typing import TYPE_CHECKING, Any, Literal +from typing import TYPE_CHECKING, Any, ClassVar, Literal from sqlalchemy import and_, case, create_engine, event, exists, func, or_, select, text from sqlalchemy.engine import AdaptedConnection, ExceptionContext @@ -190,10 +190,16 @@ class SQLiteMemory(MemoryInterface, metaclass=Singleton): Failed disconnects and session rollbacks preserve the original cancellation and expose cleanup failures as its cause. + File-backed databases wait up to 30 seconds for locks held by other connections + before failing with ``database is locked``. + Note: this is replacing the old DuckDB implementation. """ DEFAULT_DB_FILE_NAME = "pyrit.db" + # Concurrent scenario writers can wait several seconds for SQLite's single write lock, + # which outlasts sqlite3's 5 second default. + _FILE_BUSY_TIMEOUT_SECONDS: ClassVar[float] = 30.0 def __init__( self, @@ -253,9 +259,12 @@ def _initialize_schema(self) -> None: def _create_async_engine(self) -> AsyncEngine: database = self._memory_uri if self.db_path == ":memory:" else str(self.db_path) - kwargs: dict[str, Any] = {"connect_args": {"factory": _CursorClosingSQLiteConnection}} + connect_args: dict[str, Any] = {"factory": _CursorClosingSQLiteConnection} + kwargs: dict[str, Any] = {"connect_args": connect_args} if self.db_path == ":memory:": kwargs["poolclass"] = StaticPool + else: + connect_args["timeout"] = self._FILE_BUSY_TIMEOUT_SECONDS engine = create_async_engine(f"sqlite+aiosqlite:///{database}", echo=self._verbose, **kwargs) self._register_analytics_lower(engine=engine.sync_engine) event.listen(engine.sync_engine, "handle_error", _cleanup_interrupted_sqlite_connection) @@ -355,6 +364,8 @@ def _create_engine(self, *, has_echo: bool) -> Engine: # the connection will be shared across threads. extra_kwargs["poolclass"] = StaticPool extra_kwargs["connect_args"] = {"check_same_thread": False} + else: + extra_kwargs["connect_args"] = {"timeout": self._FILE_BUSY_TIMEOUT_SECONDS} database = self._memory_uri if self.db_path == ":memory:" else str(self.db_path) engine = create_engine(f"sqlite:///{database}", echo=has_echo, **extra_kwargs) diff --git a/tests/unit/memory/test_conversation_branching.py b/tests/unit/memory/test_conversation_branching.py index d0cfe7360c..d201e4319e 100644 --- a/tests/unit/memory/test_conversation_branching.py +++ b/tests/unit/memory/test_conversation_branching.py @@ -12,10 +12,10 @@ from unittest.mock import MagicMock import pytest -from sqlalchemy import create_engine, event, select +from sqlalchemy import event, select from sqlalchemy.dialects import mssql, sqlite from sqlalchemy.exc import IntegrityError, OperationalError -from sqlalchemy.orm import Session, sessionmaker +from sqlalchemy.orm import Session from pyrit.memory import MemoryInterface, SQLiteMemory from pyrit.memory.memory_models import ( @@ -26,7 +26,6 @@ PromptConverterIdentifierEntry, TargetIdentifierEntry, ) -from pyrit.memory.memory_session import MemorySession from pyrit.models import ( AttackResult, Conversation, @@ -45,18 +44,6 @@ async def file_memory(*, tmp_path: Path) -> AsyncGenerator[SQLiteMemory, None]: memory = SQLiteMemory.__new__(SQLiteMemory) memory.__init__(db_path=tmp_path / "branching.db", skip_schema_migration=True) - memory.engine.dispose() - memory.engine = create_engine(memory.engine.url, connect_args={"timeout": 30}) - memory.SessionFactory = sessionmaker(bind=memory.engine, class_=MemorySession) - - def configure_busy_timeout(*, dbapi_connection: Any, **kwargs: Any) -> None: - cursor = dbapi_connection.cursor() - try: - cursor.execute("PRAGMA busy_timeout=30000") - finally: - cursor.close() - - event.listen(memory._get_async_engine().sync_engine, "connect", configure_busy_timeout, named=True) try: Base.metadata.create_all(memory.engine) yield memory diff --git a/tests/unit/memory/test_sqlite_memory.py b/tests/unit/memory/test_sqlite_memory.py index b5798f473a..038948084b 100644 --- a/tests/unit/memory/test_sqlite_memory.py +++ b/tests/unit/memory/test_sqlite_memory.py @@ -1,16 +1,19 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +import asyncio import gc import io import logging import os +import sqlite3 import tempfile import threading import uuid from collections.abc import Callable, Sequence from contextlib import closing from pathlib import Path +from typing import Any from unittest.mock import MagicMock import pytest @@ -24,7 +27,7 @@ from pyrit.converter.base64_converter import Base64Converter from pyrit.memory.alembic.versions.ab8f2c1a9d07_pre_alembic_release_schema import INITIAL_METADATA from pyrit.memory.memory_embedding import MemoryEmbedding -from pyrit.memory.memory_models import EmbeddingDataEntry, PromptMemoryEntry +from pyrit.memory.memory_models import Base, EmbeddingDataEntry, PromptMemoryEntry from pyrit.memory.migration import run_schema_migrations from pyrit.memory.sqlite_memory import SQLiteMemory from pyrit.memory.storage.serializers import set_message_piece_sha256_async @@ -35,11 +38,14 @@ EmbeddingUsageInformation, Message, MessagePiece, + ScenarioRunState, SeedPrompt, + TargetIdentifier, flatten_to_message_pieces, ) from pyrit.prompt_target.text_target import TextTarget -from unit.mocks import get_sample_conversation_entries +from unit.async_utils import settle_tasks_async, wait_for_completion_async, wait_for_event_async +from unit.mocks import get_sample_conversation_entries, make_scenario_result @pytest.fixture @@ -1221,6 +1227,84 @@ async def test_file_backed_database_is_not_serialized_async(isolated_memory_fact await memory.dispose_engine_async() +async def test_file_backed_sessions_use_30_second_busy_timeout_async( + isolated_memory_factory: Callable[..., SQLiteMemory], tmp_path: Path +) -> None: + memory = isolated_memory_factory(db_path=tmp_path / "busy-timeout.db", skip_schema_migration=True) + + with closing(memory._get_sync_session()) as session: + assert session.execute(text("PRAGMA busy_timeout")).scalar_one() == 30_000 + async with await memory.get_session_async() as session: + assert (await session.execute(text("PRAGMA busy_timeout"))).scalar_one() == 30_000 + + +async def test_file_backed_scenario_writes_wait_for_another_connections_lock_async( + isolated_memory_factory: Callable[..., SQLiteMemory], tmp_path: Path +) -> None: + """Concurrent scenario writes wait in SQLite's busy handler rather than failing with ``database is locked``.""" + database = tmp_path / "contention.db" + memory = isolated_memory_factory(db_path=database, skip_schema_migration=True) + Base.metadata.create_all(memory.engine) + scenario_target = TargetIdentifier(class_name="ScenarioTarget", class_module="tests.unit.memory") + existing = make_scenario_result(objective_target_identifier=scenario_target, attack_results={}) + await memory.add_scenario_results_to_memory_async(scenario_results=[existing]) + inserted = make_scenario_result(objective_target_identifier=scenario_target, attack_results={}) + conversation = Conversation( + conversation_id=str(uuid.uuid4()), + target_identifier=TargetIdentifier(class_name="ContendedTarget", class_module="tests.unit.memory"), + ) + loop = asyncio.get_running_loop() + writing_connections: set[int] = set() + all_writers_waiting = asyncio.Event() + + def record_write(connection_key: int) -> None: + writing_connections.add(connection_key) + if len(writing_connections) == 3: + all_writers_waiting.set() + + def install_trace(*, dbapi_connection: Any, connection_record: Any, **kwargs: Any) -> None: + connection_key = id(connection_record) + + def trace_statement(statement: str) -> None: + if statement.startswith(("INSERT", "UPDATE")): + loop.call_soon_threadsafe(record_write, connection_key) + + dbapi_connection.run_async(lambda driver: driver.set_trace_callback(trace_statement)) + + engine = memory._get_async_engine().sync_engine + writers: list[asyncio.Task[None]] = [] + with closing(sqlite3.connect(database, isolation_level=None)) as holder: + holder.execute("BEGIN IMMEDIATE") + event.listen(engine, "checkout", install_trace, named=True) + try: + writers = [ + asyncio.create_task(memory.add_scenario_results_to_memory_async(scenario_results=[inserted])), + asyncio.create_task(memory.add_conversation_to_memory_async(conversation=conversation)), + asyncio.create_task( + memory.update_scenario_run_state_async( + scenario_result_id=str(existing.id), scenario_run_state=ScenarioRunState.COMPLETED + ) + ), + ] + await wait_for_event_async(event=all_writers_waiting, tasks=writers) + assert not any(writer.done() for writer in writers) + holder.execute("ROLLBACK") + await wait_for_completion_async(future=asyncio.gather(*writers)) + finally: + event.remove(engine, "checkout", install_trace) + if holder.in_transaction: + holder.execute("ROLLBACK") + await settle_tasks_async(*writers) + + [updated] = await memory.get_scenario_results_async(scenario_result_ids=[str(existing.id)]) + assert updated.scenario_run_state == ScenarioRunState.COMPLETED + stored_results = await memory.get_scenario_results_async(scenario_result_ids=[str(inserted.id)]) + assert [result.id for result in stored_results] == [inserted.id] + stored_conversation = await memory.get_conversation_metadata_async(conversation_id=conversation.conversation_id) + assert stored_conversation is not None + assert stored_conversation.target_identifier == conversation.target_identifier + + async def test_seed_list_substring_filters_after_reopening_database_async( isolated_memory_factory: Callable[..., SQLiteMemory], tmp_path: Path ) -> None: