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
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:
@@ -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
@@ -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
@@ -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.")
|
|
||||||
|
|||||||
@@ -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(
|
||||||
|
|||||||
+21
-80
@@ -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,15 +196,13 @@ 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:
|
|
||||||
row = await s2.fetch_one("SELECT name FROM items WHERE id = 1")
|
|
||||||
assert row is not None
|
assert row is not None
|
||||||
assert row["name"] == "a"
|
assert row["name"] == "a"
|
||||||
|
|
||||||
@@ -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
|
||||||
|
|||||||
Reference in New Issue
Block a user