Simplified storage layer from connection pool to single lazy connection.
CI / Formatting (push) Failing after 19s
CI / Linting (push) Successful in 18s
CI / Tests (Python 3.13) (push) Has been cancelled
CI / Tests (Python 3.14) (push) Has been cancelled
CI / Type Checking (push) Has been cancelled
CI / Spelling (push) Has been cancelled
CI / Tests (Python 3.12) (push) Has been cancelled

This commit is contained in:
2026-04-04 18:39:25 -04:00
parent cb1573706e
commit 09af238c8b
6 changed files with 63 additions and 199 deletions
-3
View File
@@ -49,9 +49,6 @@ owlbot:
# logs go to stdout only. # logs go to stdout only.
#log_dir: "logs" #log_dir: "logs"
# Maximum number of SQLite connections per module in the connection pool.
# Connections are created lazily as needed, up to this limit.
pool_size: 4
# Per-module configuration. Each key is a module name (the filename without # Per-module configuration. Each key is a module name (the filename without
# .py for single-file modules, or the directory name for package modules). # .py for single-file modules, or the directory name for package modules).
+1 -1
Submodule docs updated: 2f7c73138a...2258c917c8
-14
View File
@@ -206,20 +206,6 @@ class Config:
type_fn=float, type_fn=float,
) )
@property
def pool_size(self) -> int:
"""Maximum number of SQLite connections per module storage pool."""
return max(
1,
self._resolve(
env="OWLBOT_POOL_SIZE",
section="owlbot",
key="pool_size",
default=4,
type_fn=int,
),
)
@property @property
def command_prefix(self) -> str: def command_prefix(self) -> str:
"""Prefix character for chat commands.""" """Prefix character for chat commands."""
+38 -96
View File
@@ -16,8 +16,6 @@
from __future__ import annotations from __future__ import annotations
import asyncio
import contextlib
import contextvars import contextvars
import logging import logging
from contextlib import asynccontextmanager from contextlib import asynccontextmanager
@@ -35,40 +33,35 @@ class StorageError(Exception):
class ModuleStorage: class ModuleStorage:
"""Module-scoped async SQLite storage backed by a lazy connection pool. """Module-scoped async SQLite storage with a single lazy connection.
Each module gets its own isolated database file. Connections are created Each module gets its own isolated database file. The connection is
lazily and pooled up to ``pool_size``. WAL mode is enabled so that created lazily on the first storage operation and reused for all
concurrent readers and a single writer can operate without "database is subsequent calls. WAL mode is enabled for crash resilience.
locked" errors.
Each storage operation (execute, fetch_one, etc.) independently acquires a Each storage operation (execute, fetch_one, etc.) auto-commits on
connection from the pool, auto-commits on success, rolls back on failure, success and rolls back on failure. No transaction state is held
and releases the connection immediately. No connection is held between between calls.
calls.
For operations that must succeed or fail together, use the transaction() For operations that must succeed or fail together, use the
context manager to group them into a single atomic unit. transaction() context manager to group them into a single atomic
unit.
""" """
def __init__(self, storage_dir: Path, module_name: str, pool_size: int = 4): def __init__(self, storage_dir: Path, module_name: str):
"""Initialize the storage API. """Initialize the storage API.
:param storage_dir: Directory where module databases are stored. :param storage_dir: Directory where module databases are stored.
:param module_name: Name of the module this storage belongs to. :param module_name: Name of the module this storage belongs to.
:param pool_size: Maximum number of pooled connections.
""" """
self._db_path = storage_dir / f"{module_name}.db" self._db_path = storage_dir / f"{module_name}.db"
self._module_name = module_name self._module_name = module_name
self._pool: asyncio.Queue[aiosqlite.Connection] = asyncio.Queue() self._conn: aiosqlite.Connection | None = None
self._pool_size = pool_size
self._all_connections: list[aiosqlite.Connection] = []
self._closed = False self._closed = False
self._txn_conn: contextvars.ContextVar[aiosqlite.Connection | None] = ( self._txn_conn: contextvars.ContextVar[aiosqlite.Connection | None] = (
contextvars.ContextVar(f"_txn_conn_{module_name}", default=None) contextvars.ContextVar(f"_txn_conn_{module_name}", default=None)
) )
self._logger = logging.getLogger(f"owlbot.modules.{module_name}.storage") self._logger = logging.getLogger(f"owlbot.modules.{module_name}.storage")
self._logger.info(f"Database initialized at: {self._db_path.absolute()}")
async def __aenter__(self) -> ModuleStorage: async def __aenter__(self) -> ModuleStorage:
"""Enter an async context manager that closes the storage on exit. """Enter an async context manager that closes the storage on exit.
@@ -91,14 +84,13 @@ class ModuleStorage:
async def transaction(self) -> AsyncIterator[ModuleStorage]: async def transaction(self) -> AsyncIterator[ModuleStorage]:
"""Context manager for explicit transaction control. """Context manager for explicit transaction control.
Use this when you need multiple operations to succeed or fail together. Use this when you need multiple operations to succeed or fail
Acquires a dedicated connection from the pool and shares it across all together. Shares the single connection across all operations
operations within the block. Commits on success, rolls back on exception. within the block, skipping per-operation auto-commit. Commits on
success, rolls back on exception.
Nesting ``transaction()`` calls is not supported and raises Nesting ``transaction()`` calls is not supported and raises
``RuntimeError``. SQLite only allows one writer at a time, so a ``RuntimeError``.
nested transaction would deadlock waiting for the outer connection's
write lock.
Example: Example:
async with ctx.storage.transaction(): async with ctx.storage.transaction():
@@ -114,7 +106,7 @@ class ModuleStorage:
raise RuntimeError( raise RuntimeError(
"transaction() cannot be nested. Already inside an active transaction." "transaction() cannot be nested. Already inside an active transaction."
) )
conn = await self._acquire() conn = await self._ensure_connection()
token = self._txn_conn.set(conn) token = self._txn_conn.set(conn)
self._logger.debug("Explicit transaction started.") self._logger.debug("Explicit transaction started.")
try: try:
@@ -127,7 +119,6 @@ class ModuleStorage:
raise raise
finally: finally:
self._txn_conn.reset(token) self._txn_conn.reset(token)
self._release(conn)
async def execute( async def execute(
self, self,
@@ -243,95 +234,46 @@ class ModuleStorage:
return None return None
return row[0] return row[0]
async def _create_connection(self) -> aiosqlite.Connection: async def _ensure_connection(self) -> aiosqlite.Connection:
"""Create a new database connection with WAL mode and foreign keys. """Return the existing connection or create one on first use."""
The connection is tracked in ``_all_connections`` immediately after
opening so that ``_close()`` can clean it up if the PRAGMA setup is
interrupted (e.g. by task cancellation). A ``busy_timeout`` is set
first so that new connections briefly wait for any lock held by a
partially-initialised peer instead of failing immediately.
"""
self._db_path.parent.mkdir(parents=True, exist_ok=True)
conn = await aiosqlite.connect(self._db_path)
# Track immediately so _close() always sees this connection.
self._all_connections.append(conn)
try:
conn.row_factory = aiosqlite.Row
await conn.execute("PRAGMA busy_timeout = 1000")
await conn.execute("PRAGMA journal_mode = WAL")
await conn.execute("PRAGMA foreign_keys = ON")
except BaseException:
# Setup interrupted (timeout, cancellation, etc.). Close the
# connection to release any write lock held by a partially-
# executed PRAGMA, then free the pool slot. shield() keeps the
# close running on the background thread even if this task is
# cancelled; suppress swallows the CancelledError from the await.
with contextlib.suppress(ValueError):
self._all_connections.remove(conn)
with contextlib.suppress(asyncio.CancelledError, Exception):
await asyncio.shield(conn.close())
raise
self._logger.debug(
"Pool connection created (%d/%d).",
len(self._all_connections),
self._pool_size,
)
return conn
async def _acquire(self) -> aiosqlite.Connection:
"""Acquire a connection from the pool, creating one if needed."""
if self._closed: if self._closed:
raise StorageError("Storage is closed") raise StorageError("Storage is closed")
if self._conn is not None:
try: return self._conn
return self._pool.get_nowait() self._db_path.parent.mkdir(parents=True, exist_ok=True)
except asyncio.QueueEmpty: conn = await aiosqlite.connect(self._db_path)
pass conn.row_factory = aiosqlite.Row
await conn.execute("PRAGMA journal_mode = WAL")
if len(self._all_connections) < self._pool_size: await conn.execute("PRAGMA foreign_keys = ON")
return await self._create_connection() self._conn = conn
self._logger.info(f"Database opened at: {self._db_path.absolute()}")
# Pool exhausted, wait for one to be returned. return conn
return await self._pool.get()
def _release(self, conn: aiosqlite.Connection) -> None:
"""Return a connection to the pool."""
self._pool.put_nowait(conn)
@asynccontextmanager @asynccontextmanager
async def _connection(self) -> AsyncIterator[aiosqlite.Connection]: async def _connection(self) -> AsyncIterator[aiosqlite.Connection]:
"""Async context manager that provides a connection. """Async context manager that provides the connection.
If already inside a ``transaction()``, yields the shared connection If already inside a ``transaction()``, yields the shared connection
without committing (the transaction block handles that). Otherwise without committing (the transaction block handles that). Otherwise
acquires a standalone connection from the pool that auto-commits on auto-commits on success and rolls back on failure.
success and rolls back on failure before being released.
""" """
existing = self._txn_conn.get() existing = self._txn_conn.get()
if existing is not None: if existing is not None:
yield existing yield existing
return return
conn = await self._acquire() conn = await self._ensure_connection()
try: try:
yield conn yield conn
await conn.commit() await conn.commit()
except BaseException: except BaseException:
await conn.rollback() await conn.rollback()
raise raise
finally:
self._release(conn)
async def _close(self) -> None: async def _close(self) -> None:
"""Close all pool connections (internal use by bot).""" """Close the connection (internal use by bot)."""
self._closed = True self._closed = True
if self._conn is not None:
for conn in self._all_connections: await self._conn.close()
await conn.close() self._conn = None
self._logger.info("Connection closed.")
while not self._pool.empty():
self._pool.get_nowait()
self._all_connections.clear()
self._logger.info("All pool connections closed.")
+1 -3
View File
@@ -320,9 +320,7 @@ class ModuleLoader:
module_templates = ModuleTemplates(module_dir, self._core_template_dir) module_templates = ModuleTemplates(module_dir, self._core_template_dir)
scoped_config = ModuleConfig(self.config, module_name) scoped_config = ModuleConfig(self.config, module_name)
storage = ModuleStorage( storage = ModuleStorage(self.config.storage_dir, module_name)
self.config.storage_dir, module_name, self.config.pool_size
)
module_commands = ModuleCommands(self.command_dispatcher, module_name) module_commands = ModuleCommands(self.command_dispatcher, module_name)
module_events = ModuleEvents(self.event_dispatcher, module_name) module_events = ModuleEvents(self.event_dispatcher, module_name)
module_routes = ModuleRoutes( module_routes = ModuleRoutes(
+23 -82
View File
@@ -39,18 +39,18 @@ async def storage_with_table(storage: ModuleStorage) -> ModuleStorage:
class TestLazyConnections: class TestLazyConnections:
"""No connections exist until the first operation.""" """No connection exists until the first operation."""
async def test_no_connections_before_use(self, storage: ModuleStorage) -> None: async def test_no_connection_before_use(self, storage: ModuleStorage) -> None:
"""Pool is empty immediately after construction.""" """Connection is None immediately after construction."""
assert len(storage._all_connections) == 0 assert storage._conn is None
async def test_first_operation_creates_one_connection( async def test_first_operation_creates_connection(
self, storage: ModuleStorage self, storage: ModuleStorage
) -> None: ) -> None:
"""The first execute creates exactly one pooled connection.""" """The first execute creates the connection."""
await storage.execute("SELECT 1") await storage.execute("SELECT 1")
assert len(storage._all_connections) == 1 assert storage._conn is not None
class TestPragmas: class TestPragmas:
@@ -196,17 +196,15 @@ class TestAutoCommit:
async def test_auto_commits_on_success( async def test_auto_commits_on_success(
self, storage_with_table: ModuleStorage self, storage_with_table: ModuleStorage
) -> None: ) -> None:
"""A successful write is visible from a second storage instance.""" """A successful write is visible on the next read."""
await storage_with_table.execute( await storage_with_table.execute(
"INSERT INTO items (name, value) VALUES (?, ?)", ("a", 1) "INSERT INTO items (name, value) VALUES (?, ?)", ("a", 1)
) )
# Open a second storage pointing at the same file to verify commit. row = await storage_with_table.fetch_one(
async with ModuleStorage( "SELECT name FROM items WHERE id = 1"
storage_with_table._db_path.parent, "test_module" )
) as s2: assert row is not None
row = await s2.fetch_one("SELECT name FROM items WHERE id = 1") assert row["name"] == "a"
assert row is not None
assert row["name"] == "a"
async def test_auto_rollback_on_error( async def test_auto_rollback_on_error(
self, storage_with_table: ModuleStorage self, storage_with_table: ModuleStorage
@@ -267,21 +265,6 @@ class TestTransaction:
async with storage_with_table.transaction(): async with storage_with_table.transaction():
pass # pragma: no cover pass # pragma: no cover
async def test_uncommitted_transaction_invisible_from_other_connection(
self, storage_with_table: ModuleStorage
) -> None:
"""Uncommitted writes inside a transaction are not visible externally."""
async with (
ModuleStorage(storage_with_table._db_path.parent, "test_module") as s2,
storage_with_table.transaction(),
):
await storage_with_table.execute(
"INSERT INTO items (name, value) VALUES (?, ?)", ("a", 1)
)
# Before commit, s2 should not see the row.
row = await s2.fetch_one("SELECT name FROM items WHERE id = 1")
assert row is None
async def test_rolls_back_on_cancellation( async def test_rolls_back_on_cancellation(
self, storage_with_table: ModuleStorage self, storage_with_table: ModuleStorage
) -> None: ) -> None:
@@ -304,48 +287,6 @@ class TestTransaction:
assert row is None assert row is None
class TestConnectionPool:
"""Connection pool reuses connections and respects pool_size."""
async def test_connections_are_reused(self, storage: ModuleStorage) -> None:
"""After release, the same connection is reused."""
await storage.execute("SELECT 1")
assert len(storage._all_connections) == 1
await storage.execute("SELECT 2")
assert len(storage._all_connections) == 1
async def test_pool_grows_up_to_pool_size(self, tmp_path: Path) -> None:
"""Concurrent transactions grow the pool up to pool_size."""
gate = asyncio.Event()
async def hold_transaction(storage: ModuleStorage) -> None:
async with storage.transaction():
gate.set()
await asyncio.sleep(0.5)
async with ModuleStorage(tmp_path, "pool_test", pool_size=2) as s:
task = asyncio.create_task(hold_transaction(s))
await gate.wait()
# First task holds one connection; acquire a second directly.
conn = await asyncio.wait_for(s._acquire(), timeout=1.0)
assert len(s._all_connections) == 2
s._release(conn)
task.cancel()
with contextlib.suppress(asyncio.CancelledError):
await task
async def test_pool_exhaustion_blocks(self, tmp_path: Path) -> None:
"""When the pool is exhausted, acquire blocks until a connection is released."""
async with (
ModuleStorage(tmp_path, "pool_block", pool_size=1) as s,
s.transaction(),
):
# Pool is exhausted (pool_size=1, one held by transaction).
# A second acquire should block, so wait_for should time out.
with pytest.raises(TimeoutError):
await asyncio.wait_for(s._acquire(), timeout=0.1)
class TestStorageError: class TestStorageError:
"""Invalid SQL raises StorageError wrapping the underlying error.""" """Invalid SQL raises StorageError wrapping the underlying error."""
@@ -412,15 +353,15 @@ class TestStorageError:
class TestClose: class TestClose:
"""_close() drains the pool and marks storage closed.""" """_close() closes the connection and marks storage closed."""
async def test_close_drains_pool(self, storage: ModuleStorage) -> None: async def test_close_releases_connection(self, storage: ModuleStorage) -> None:
"""After close, all connections are removed.""" """After close, the connection is None."""
await storage.execute("SELECT 1") await storage.execute("SELECT 1")
assert len(storage._all_connections) == 1 assert storage._conn is not None
await storage._close() await storage._close()
assert len(storage._all_connections) == 0 assert storage._closed
assert storage._closed is True assert storage._conn is None
async def test_double_close_is_safe(self, storage: ModuleStorage) -> None: async def test_double_close_is_safe(self, storage: ModuleStorage) -> None:
"""Calling _close() twice does not raise.""" """Calling _close() twice does not raise."""
@@ -429,9 +370,9 @@ class TestClose:
await storage._close() await storage._close()
async def test_context_manager_closes_on_exit(self, tmp_path: Path) -> None: async def test_context_manager_closes_on_exit(self, tmp_path: Path) -> None:
"""Exiting the async with block closes all connections.""" """Exiting the async with block closes the connection."""
async with ModuleStorage(tmp_path, "ctx_test") as s: async with ModuleStorage(tmp_path, "ctx_test") as s:
await s.execute("SELECT 1") await s.execute("SELECT 1")
assert len(s._all_connections) == 1 assert s._conn is not None
assert len(s._all_connections) == 0 assert s._closed
assert s._closed is True assert s._conn is None