Owlbot/owlbot/api/storage.py

# Copyright 2026 Logan Fick
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""SQLite storage API for Owlbot modules."""

from __future__ import annotations

import asyncio
import contextlib
import contextvars
import logging
from contextlib import asynccontextmanager
from pathlib import Path
from typing import TYPE_CHECKING, Any, Self

import aiosqlite

if TYPE_CHECKING:
    from collections.abc import AsyncIterator


class StorageError(Exception):
    """Raised when a storage operation fails."""


class ModuleStorage:
    """Module-scoped async SQLite storage with a single lazy connection.

    Each module gets its own isolated database file. The connection is
    created lazily on the first storage operation and reused for all
    subsequent calls. WAL mode is enabled for crash resilience.

    Pass ``None`` as ``storage_dir`` to use an in-memory database
    instead of a file, which is useful for testing.

    Each storage operation (execute, fetch_one, etc.) auto-commits on
    success and rolls back on failure. No transaction state is held
    between calls.

    For operations that must succeed or fail together, use the
    transaction() context manager to group them into a single atomic
    unit.
    """

    def __init__(self, storage_dir: Path | None, module_name: str) -> None:
        """Initialize the storage API.

        :param storage_dir: Directory where module databases are stored.
            Pass ``None`` to use an in-memory database (useful for tests).
        :param module_name: Name of the module this storage belongs to.
        """
        self._db_path: Path | str = (
            storage_dir / f"{module_name}.db" if storage_dir is not None else ":memory:"
        )
        self._module_name = module_name
        self._conn: aiosqlite.Connection | None = None
        self._closed = False
        self._lock = asyncio.Lock()
        self._txn_conn: contextvars.ContextVar[aiosqlite.Connection | None] = (
            contextvars.ContextVar(f"_txn_conn_{module_name}", default=None)
        )
        self._logger = logging.getLogger(f"owlbot.modules.{module_name}.storage")

    async def __aenter__(self) -> Self:
        """Enter an async context manager that closes the storage on exit.

        Enables ``async with ModuleStorage(...) as s:`` for scoped usage
        (e.g. test fixtures).

        In production, ``ModuleLoader`` manages storage lifecycle with explicit
        ``_close()`` calls because the storage is created during module load
        but closed in separate code paths (unload, load failure, setup failure),
        so there is no single scope an ``async with`` block could wrap.
        """
        return self

    async def __aexit__(self, *exc_info: object) -> None:
        """Close the storage when leaving the ``async with`` block."""
        await self._close()

    @asynccontextmanager
    async def transaction(self) -> AsyncIterator[ModuleStorage]:
        """Context manager for explicit transaction control.

        Use this when you need multiple operations to succeed or fail
        together. Shares the single connection across all operations
        within the block, skipping per-operation auto-commit. Commits on
        success, rolls back on exception.

        Nesting ``transaction()`` calls is not supported and raises
        ``RuntimeError``.

        Example:
            async with ctx.storage.transaction():
                await ctx.storage.execute("UPDATE scores SET score = score - ?", (10,))
                await ctx.storage.execute("UPDATE scores SET score = score + ?", (10,))
                # Both committed together, or both rolled back on error

        :return: This ModuleStorage instance.
        :raises RuntimeError: If called while already inside a transaction.

        """
        if self._txn_conn.get() is not None:
            raise RuntimeError(
                "transaction() cannot be nested. Already inside an active transaction."
            )
        async with self._lock:
            conn = await self._ensure_connection()
            token = self._txn_conn.set(conn)
            self._logger.debug("Explicit transaction started.")
            try:
                yield self
                await conn.commit()
                self._logger.debug("Transaction committed.")
            except BaseException:
                await conn.rollback()
                self._logger.debug("Transaction rolled back.")
                raise
            finally:
                self._txn_conn.reset(token)

    async def execute(
        self,
        sql: str,
        parameters: tuple[Any, ...] | dict[str, Any] = (),
    ) -> aiosqlite.Cursor:
        """Execute a SQL statement.

        :param sql: SQL statement (use ? or :name for parameters).
        :param parameters: Query parameters (tuple for ?, dict for :name).
        :return: Cursor with lastrowid and rowcount.
        :raises StorageError: If execution fails.
        """
        async with self._connection() as conn:
            self._logger.debug("Execute: %s", sql)
            try:
                return await conn.execute(sql, parameters)
            except aiosqlite.Error as e:
                self._logger.exception("SQL error.")
                raise StorageError(f"SQL execution failed: {e}") from e

    async def execute_many(
        self,
        sql: str,
        parameters: list[tuple[Any, ...]] | list[dict[str, Any]],
    ) -> aiosqlite.Cursor:
        """Execute a SQL statement with multiple parameter sets.

        Useful for batch inserts/updates.

        :param sql: SQL statement.
        :param parameters: List of parameter tuples/dicts.
        :return: Cursor with rowcount.
        :raises StorageError: If execution fails.
        """
        async with self._connection() as conn:
            self._logger.debug(
                "Execute many (%d rows): %s",
                len(parameters),
                sql,
            )
            try:
                return await conn.executemany(sql, parameters)
            except aiosqlite.Error as e:
                self._logger.exception("SQL error in executemany.")
                raise StorageError(f"SQL execution failed: {e}") from e

    async def fetch_one(
        self,
        sql: str,
        parameters: tuple[Any, ...] | dict[str, Any] = (),
    ) -> aiosqlite.Row | None:
        """Execute a query and fetch one row.

        :param sql: SELECT statement.
        :param parameters: Query parameters.
        :return: Row as a sqlite3.Row (supports both index and key access),
            or None if no results.
        :raises StorageError: If execution fails.
        """
        async with self._connection() as conn:
            self._logger.debug("Fetch one: %s", sql)
            try:
                cursor = await conn.execute(sql, parameters)
                return await cursor.fetchone()
            except aiosqlite.Error as e:
                self._logger.exception("SQL error.")
                raise StorageError(f"SQL fetch failed: {e}") from e

    async def fetch_all(
        self,
        sql: str,
        parameters: tuple[Any, ...] | dict[str, Any] = (),
    ) -> list[aiosqlite.Row]:
        """Execute a query and fetch all rows.

        :param sql: SELECT statement.
        :param parameters: Query parameters.
        :return: List of rows as sqlite3.Row objects
            (support both index and key access).
        :raises StorageError: If execution fails.
        """
        async with self._connection() as conn:
            self._logger.debug("Fetch all: %s", sql)
            try:
                cursor = await conn.execute(sql, parameters)
                return list(await cursor.fetchall())
            except aiosqlite.Error as e:
                self._logger.exception("SQL error.")
                raise StorageError(f"SQL fetch failed: {e}") from e

    async def fetch_value(
        self,
        sql: str,
        parameters: tuple[Any, ...] | dict[str, Any] = (),
    ) -> Any | None:
        """Execute a query and fetch a single value.

        :param sql: SELECT statement returning one column.
        :param parameters: Query parameters.
        :return: The value, or None if no results.
        :raises StorageError: If execution fails.
        """
        row = await self.fetch_one(sql, parameters)
        if row is None:
            return None
        return row[0]

    async def _ensure_connection(self) -> aiosqlite.Connection:
        """Return the existing connection or create one on first use."""
        if self._closed:
            raise StorageError("Storage is closed")
        if self._conn is not None:
            return self._conn
        if isinstance(self._db_path, Path):
            self._db_path.parent.mkdir(parents=True, exist_ok=True)
        conn = await aiosqlite.connect(self._db_path)
        try:
            conn.row_factory = aiosqlite.Row
            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. shield() keeps the close running on the
            # background thread even if this task is cancelled.
            with contextlib.suppress(asyncio.CancelledError, Exception):
                await asyncio.shield(conn.close())
            raise
        self._conn = conn
        self._logger.info("Database opened at: %s", self._db_path)
        return conn

    @asynccontextmanager
    async def _connection(self) -> AsyncIterator[aiosqlite.Connection]:
        """Async context manager that provides the connection.

        If already inside a ``transaction()``, yields the shared connection
        without committing (the transaction block handles that). Otherwise
        auto-commits on success and rolls back on failure.
        """
        existing = self._txn_conn.get()
        if existing is not None:
            yield existing
            return

        async with self._lock:
            conn = await self._ensure_connection()
            try:
                yield conn
                await conn.commit()
            except BaseException:
                await conn.rollback()
                raise

    async def _close(self) -> None:
        """Close the connection (internal use by bot)."""
        async with self._lock:
            self._closed = True
            if self._conn is not None:
                await self._conn.close()
                self._conn = None
        self._logger.info("Connection closed.")