kstlib 0.0.1a0__py3-none-any.whl → 1.0.0__py3-none-any.whl

kstlib/db/database.py

@@ -0,0 +1,367 @@
+"""Async database wrapper for SQLite/SQLCipher.
+
+Provides a high-level async interface for database operations with:
+- Connection pooling
+- Automatic retry on transient failures
+- SQLCipher encryption support
+- Transaction management
+- Query helpers
+"""
+
+from __future__ import annotations
+
+import logging
+from contextlib import asynccontextmanager
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, cast
+
+from typing_extensions import Self
+
+from kstlib.db.exceptions import TransactionError
+from kstlib.db.pool import ConnectionPool, PoolStats
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncGenerator, Sequence
+    from pathlib import Path
+
+    import aiosqlite
+
+log = logging.getLogger(__name__)
+
+
+@dataclass
+class AsyncDatabase:
+    """Async database interface for SQLite/SQLCipher.
+
+    Provides connection pooling, encryption, and query helpers
+    for async database operations.
+
+    Args:
+        path: Path to database file (or ":memory:" for in-memory).
+        cipher_key: Direct encryption key for SQLCipher.
+        cipher_env: Environment variable containing cipher key.
+        cipher_sops: Path to SOPS file containing cipher key.
+        cipher_sops_key: Key name in SOPS file (default: "db_key").
+        pool_min: Minimum pool connections.
+        pool_max: Maximum pool connections.
+        pool_timeout: Acquire timeout in seconds.
+        max_retries: Retry attempts on failure.
+        retry_delay: Delay between retries.
+
+    Examples:
+        Basic usage:
+
+        >>> db = AsyncDatabase(":memory:")
+        >>> db.path
+        ':memory:'
+
+        With encryption:
+
+        >>> db = AsyncDatabase("app.db", cipher_key="secret")  # doctest: +SKIP
+
+        With SOPS:
+
+        >>> db = AsyncDatabase("app.db", cipher_sops="secrets.yml")  # doctest: +SKIP
+    """
+
+    path: str | Path
+    cipher_key: str | None = None
+    cipher_env: str | None = None
+    cipher_sops: str | Path | None = None
+    cipher_sops_key: str = "db_key"
+    pool_min: int = 1
+    pool_max: int = 10
+    pool_timeout: float = 30.0
+    max_retries: int = 3
+    retry_delay: float = 0.5
+
+    _pool: ConnectionPool | None = field(default=None, repr=False)
+    _resolved_key: str | None = field(default=None, repr=False)
+
+    def __post_init__(self) -> None:
+        """Resolve cipher key and apply config defaults with hard limits."""
+        from kstlib.limits import get_db_limits
+
+        self.path = str(self.path)
+
+        # Load config defaults for any unset pool/retry params
+        limits = get_db_limits()
+
+        # Apply config defaults if using dataclass defaults (sentinel check)
+        # Use object.__setattr__ since dataclass fields are set
+        if self.pool_min == 1:
+            object.__setattr__(self, "pool_min", limits.pool_min_size)
+        if self.pool_max == 10:
+            object.__setattr__(self, "pool_max", limits.pool_max_size)
+        if self.pool_timeout == 30.0:
+            object.__setattr__(self, "pool_timeout", limits.pool_acquire_timeout)
+        if self.max_retries == 3:
+            object.__setattr__(self, "max_retries", limits.max_retries)
+        if self.retry_delay == 0.5:
+            object.__setattr__(self, "retry_delay", limits.retry_delay)
+
+        # Resolve encryption key if any source provided
+        if self.cipher_key or self.cipher_env or self.cipher_sops:
+            from kstlib.db.cipher import resolve_cipher_key
+
+            self._resolved_key = resolve_cipher_key(
+                passphrase=self.cipher_key,
+                env_var=self.cipher_env,
+                sops_path=self.cipher_sops,
+                sops_key=self.cipher_sops_key,
+            )
+
+    def _ensure_pool(self) -> ConnectionPool:
+        """Ensure connection pool is initialized."""
+        if self._pool is None:
+            # path is already converted to str in __post_init__
+            db_path = self.path if isinstance(self.path, str) else str(self.path)
+            self._pool = ConnectionPool(
+                db_path=db_path,
+                min_size=self.pool_min,
+                max_size=self.pool_max,
+                acquire_timeout=self.pool_timeout,
+                max_retries=self.max_retries,
+                retry_delay=self.retry_delay,
+                cipher_key=self._resolved_key,
+            )
+        return self._pool
+
+    async def connect(self) -> None:
+        """Initialize the connection pool.
+
+        Called automatically on first operation, but can be
+        called explicitly for eager initialization.
+        """
+        pool = self._ensure_pool()
+        await pool._init_pool()
+        log.info("Database connected: %s", self.path)
+
+    async def close(self) -> None:
+        """Close all connections and shutdown the pool."""
+        if self._pool:
+            await self._pool.close()
+            self._pool = None
+        # Scrub resolved key from memory
+        self._resolved_key = None
+        log.info("Database closed: %s", self.path)
+
+    async def __aenter__(self) -> Self:
+        """Async context manager entry."""
+        await self.connect()
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: object,
+    ) -> None:
+        """Async context manager exit."""
+        await self.close()
+
+    @asynccontextmanager
+    async def connection(self) -> AsyncGenerator[aiosqlite.Connection, None]:
+        """Get a connection from the pool.
+
+        Yields:
+            Database connection.
+
+        Examples:
+            >>> async with db.connection() as conn:  # doctest: +SKIP
+            ...     await conn.execute("SELECT 1")
+        """
+        pool = self._ensure_pool()
+        async with pool.connection() as conn:
+            yield conn
+
+    @asynccontextmanager
+    async def transaction(self) -> AsyncGenerator[aiosqlite.Connection, None]:
+        """Execute operations within a transaction.
+
+        Automatically commits on success, rolls back on error.
+
+        Yields:
+            Database connection within transaction.
+
+        Raises:
+            TransactionError: If transaction fails.
+
+        Examples:
+            >>> async with db.transaction() as conn:  # doctest: +SKIP
+            ...     await conn.execute("INSERT INTO users VALUES (?)", ("alice",))
+            ...     await conn.execute("INSERT INTO users VALUES (?)", ("bob",))
+        """
+        pool = self._ensure_pool()
+        conn = await pool.acquire()
+        try:
+            await conn.execute("BEGIN")
+            yield conn
+            await conn.commit()
+        except Exception as e:
+            try:
+                await conn.rollback()
+            except Exception:
+                # Rollback may fail on closed connection - intentional silent catch
+                log.debug("Rollback failed (connection may be closed)", exc_info=True)
+            raise TransactionError(f"Transaction failed: {e}") from e
+        finally:
+            await pool.release(conn)
+
+    async def execute(
+        self,
+        sql: str,
+        parameters: Sequence[Any] | None = None,
+    ) -> aiosqlite.Cursor:
+        """Execute a single SQL statement.
+
+        Args:
+            sql: SQL statement to execute.
+            parameters: Query parameters.
+
+        Returns:
+            Cursor with results.
+
+        Examples:
+            >>> await db.execute("CREATE TABLE test (id INTEGER)")  # doctest: +SKIP
+        """
+        pool = self._ensure_pool()
+        async with pool.connection() as conn:
+            if parameters:
+                return await conn.execute(sql, parameters)
+            return await conn.execute(sql)
+
+    async def executemany(
+        self,
+        sql: str,
+        parameters: Sequence[Sequence[Any]],
+    ) -> aiosqlite.Cursor:
+        """Execute SQL statement for multiple parameter sets.
+
+        Args:
+            sql: SQL statement to execute.
+            parameters: Sequence of parameter tuples.
+
+        Returns:
+            Cursor with results.
+
+        Examples:
+            >>> await db.executemany(  # doctest: +SKIP
+            ...     "INSERT INTO test VALUES (?)",
+            ...     [(1,), (2,), (3,)]
+            ... )
+        """
+        pool = self._ensure_pool()
+        async with pool.connection() as conn:
+            return await conn.executemany(sql, parameters)
+
+    async def fetch_one(
+        self,
+        sql: str,
+        parameters: Sequence[Any] | None = None,
+    ) -> tuple[Any, ...] | None:
+        """Fetch a single row.
+
+        Args:
+            sql: SQL query.
+            parameters: Query parameters.
+
+        Returns:
+            Row tuple or None if no results.
+
+        Examples:
+            >>> row = await db.fetch_one("SELECT * FROM test WHERE id=?", (1,))  # doctest: +SKIP
+        """
+        pool = self._ensure_pool()
+        async with pool.connection() as conn:
+            if parameters:
+                cursor = await conn.execute(sql, parameters)
+            else:
+                cursor = await conn.execute(sql)
+            row = await cursor.fetchone()
+            return cast("tuple[Any, ...] | None", row)
+
+    async def fetch_all(
+        self,
+        sql: str,
+        parameters: Sequence[Any] | None = None,
+    ) -> list[tuple[Any, ...]]:
+        """Fetch all rows.
+
+        Args:
+            sql: SQL query.
+            parameters: Query parameters.
+
+        Returns:
+            List of row tuples.
+
+        Examples:
+            >>> rows = await db.fetch_all("SELECT * FROM test")  # doctest: +SKIP
+        """
+        pool = self._ensure_pool()
+        async with pool.connection() as conn:
+            if parameters:
+                cursor = await conn.execute(sql, parameters)
+            else:
+                cursor = await conn.execute(sql)
+            rows = await cursor.fetchall()
+            return cast("list[tuple[Any, ...]]", rows)
+
+    async def fetch_value(
+        self,
+        sql: str,
+        parameters: Sequence[Any] | None = None,
+    ) -> Any:
+        """Fetch a single value (first column of first row).
+
+        Args:
+            sql: SQL query.
+            parameters: Query parameters.
+
+        Returns:
+            Single value or None.
+
+        Examples:
+            >>> count = await db.fetch_value("SELECT count(*) FROM test")  # doctest: +SKIP
+        """
+        row = await self.fetch_one(sql, parameters)
+        return row[0] if row else None
+
+    async def table_exists(self, table_name: str) -> bool:
+        """Check if a table exists.
+
+        Args:
+            table_name: Name of the table.
+
+        Returns:
+            True if table exists.
+
+        Examples:
+            >>> await db.table_exists("users")  # doctest: +SKIP
+            False
+        """
+        sql = "SELECT count(*) FROM sqlite_master WHERE type='table' AND name=?"
+        count = await self.fetch_value(sql, (table_name,))
+        return bool(count)
+
+    @property
+    def stats(self) -> PoolStats:
+        """Get connection pool statistics."""
+        if self._pool:
+            return self._pool.stats
+        return PoolStats()
+
+    @property
+    def is_encrypted(self) -> bool:
+        """Whether database is configured for encryption."""
+        return self._resolved_key is not None
+
+    @property
+    def pool_size(self) -> int:
+        """Current number of connections in pool."""
+        if self._pool:
+            return self._pool.size
+        return 0
+
+
+__all__ = ["AsyncDatabase"]

kstlib/db/exceptions.py

@@ -0,0 +1,25 @@
+"""Database module exceptions."""
+
+from __future__ import annotations
+
+from kstlib.config.exceptions import KstlibError
+
+
+class DatabaseError(KstlibError):
+    """Base exception for database operations."""
+
+
+class DatabaseConnectionError(DatabaseError):
+    """Failed to establish database connection."""
+
+
+class EncryptionError(DatabaseError):
+    """Failed to decrypt or access encrypted database."""
+
+
+class PoolExhaustedError(DatabaseError):
+    """Connection pool exhausted, no connections available."""
+
+
+class TransactionError(DatabaseError):
+    """Transaction operation failed."""

kstlib/db/pool.py

@@ -0,0 +1,302 @@
+"""Connection pool with retry support for async SQLite.
+
+Provides connection pooling with:
+- Configurable pool size
+- Connection health checks
+- Automatic retry on transient failures
+- Graceful shutdown
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from contextlib import asynccontextmanager
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any
+
+from kstlib.db.exceptions import DatabaseConnectionError, PoolExhaustedError
+from kstlib.limits import (
+    HARD_MAX_DB_MAX_RETRIES,
+    HARD_MAX_DB_RETRY_DELAY,
+    HARD_MAX_POOL_ACQUIRE_TIMEOUT,
+    HARD_MAX_POOL_MAX_SIZE,
+    HARD_MAX_POOL_MIN_SIZE,
+    HARD_MIN_DB_MAX_RETRIES,
+    HARD_MIN_DB_RETRY_DELAY,
+    HARD_MIN_POOL_ACQUIRE_TIMEOUT,
+    HARD_MIN_POOL_MAX_SIZE,
+    HARD_MIN_POOL_MIN_SIZE,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncGenerator
+
+    import aiosqlite
+
+log = logging.getLogger(__name__)
+
+
+@dataclass
+class PoolStats:
+    """Statistics for connection pool monitoring.
+
+    Attributes:
+        total_connections: Total connections created.
+        active_connections: Currently in-use connections.
+        idle_connections: Available connections in pool.
+        total_acquired: Total acquire operations.
+        total_released: Total release operations.
+        total_timeouts: Acquire operations that timed out.
+        total_errors: Connection errors encountered.
+
+    Examples:
+        >>> stats = PoolStats()
+        >>> stats.total_connections
+        0
+    """
+
+    total_connections: int = 0
+    active_connections: int = 0
+    idle_connections: int = 0
+    total_acquired: int = 0
+    total_released: int = 0
+    total_timeouts: int = 0
+    total_errors: int = 0
+
+
+@dataclass
+class ConnectionPool:
+    """Async connection pool for SQLite/SQLCipher databases.
+
+    Manages a pool of database connections with health checks
+    and automatic retry on failures.
+
+    Args:
+        db_path: Path to database file.
+        min_size: Minimum connections to maintain.
+        max_size: Maximum connections allowed.
+        acquire_timeout: Timeout for acquiring connection.
+        max_retries: Retry attempts on failure.
+        retry_delay: Delay between retries.
+        cipher_key: Optional encryption key for SQLCipher.
+        on_connect: Callback after connection established.
+
+    Examples:
+        >>> pool = ConnectionPool(":memory:", min_size=1, max_size=5)
+        >>> pool.max_size
+        5
+    """
+
+    db_path: str
+    min_size: int = 1
+    max_size: int = 10
+    acquire_timeout: float = 30.0
+    max_retries: int = 3
+    retry_delay: float = 0.5
+    cipher_key: str | None = None
+    on_connect: Any | None = None  # Callable[[aiosqlite.Connection], Awaitable[None]]
+
+    _pool: asyncio.Queue[aiosqlite.Connection] = field(default_factory=lambda: asyncio.Queue(), repr=False)
+    _connections: set[aiosqlite.Connection] = field(default_factory=set, repr=False)
+    _lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False)
+    _closed: bool = field(default=False, repr=False)
+    _stats: PoolStats = field(default_factory=PoolStats, repr=False)
+
+    def __post_init__(self) -> None:
+        """Validate and clamp configuration values to hard limits."""
+        # Clamp pool sizes
+        self.min_size = max(HARD_MIN_POOL_MIN_SIZE, min(HARD_MAX_POOL_MIN_SIZE, self.min_size))
+        self.max_size = max(HARD_MIN_POOL_MAX_SIZE, min(HARD_MAX_POOL_MAX_SIZE, self.max_size))
+
+        # Ensure min_size <= max_size
+        self.min_size = min(self.min_size, self.max_size)
+
+        # Clamp timeouts and delays
+        self.acquire_timeout = max(
+            HARD_MIN_POOL_ACQUIRE_TIMEOUT, min(HARD_MAX_POOL_ACQUIRE_TIMEOUT, self.acquire_timeout)
+        )
+        self.max_retries = max(HARD_MIN_DB_MAX_RETRIES, min(HARD_MAX_DB_MAX_RETRIES, self.max_retries))
+        self.retry_delay = max(HARD_MIN_DB_RETRY_DELAY, min(HARD_MAX_DB_RETRY_DELAY, self.retry_delay))
+
+    async def _create_connection(self) -> aiosqlite.Connection:
+        """Create a new database connection.
+
+        Uses aiosqlcipher for encrypted connections (when cipher_key is set),
+        or standard aiosqlite for unencrypted connections.
+        """
+        if self.cipher_key:
+            # Use SQLCipher for encrypted database
+            from kstlib.db.aiosqlcipher import connect as aiosqlcipher_connect
+
+            conn = await aiosqlcipher_connect(self.db_path, cipher_key=self.cipher_key)
+        else:
+            # Use standard sqlite3 for unencrypted database
+            import aiosqlite
+
+            conn = await aiosqlite.connect(self.db_path)
+
+        # Enable WAL mode for better concurrency (works with both)
+        await conn.execute("PRAGMA journal_mode=WAL")
+        await conn.execute("PRAGMA foreign_keys=ON")
+
+        # Call custom on_connect handler
+        if self.on_connect:
+            await self.on_connect(conn)
+
+        self._stats.total_connections += 1
+        log.debug("Created new database connection (total: %d)", self._stats.total_connections)
+        return conn
+
+    async def _init_pool(self) -> None:
+        """Initialize the connection pool with min_size connections."""
+        async with self._lock:
+            for _ in range(self.min_size):
+                conn = await self._create_connection()
+                self._connections.add(conn)
+                await self._pool.put(conn)
+            self._stats.idle_connections = self.min_size
+
+    async def acquire(self) -> aiosqlite.Connection:
+        """Acquire a connection from the pool.
+
+        Returns:
+            Database connection.
+
+        Raises:
+            PoolExhaustedError: If no connection available within timeout.
+            DatabaseConnectionError: If connection creation fails after retries.
+        """
+        if self._closed:
+            raise DatabaseConnectionError("Pool is closed")
+
+        # Initialize pool on first acquire
+        if not self._connections:
+            await self._init_pool()
+
+        for attempt in range(self.max_retries):
+            try:
+                # Try to get from pool with timeout
+                try:
+                    conn = await asyncio.wait_for(self._pool.get(), timeout=self.acquire_timeout)
+                    self._stats.idle_connections -= 1
+                except asyncio.TimeoutError:
+                    # Pool empty, try to create new if under max
+                    async with self._lock:
+                        if len(self._connections) < self.max_size:
+                            conn = await self._create_connection()
+                            self._connections.add(conn)
+                        else:
+                            self._stats.total_timeouts += 1
+                            raise PoolExhaustedError(
+                                f"Pool exhausted (max={self.max_size}), timeout after {self.acquire_timeout}s"
+                            ) from None
+
+                # Verify connection is alive
+                try:
+                    await conn.execute("SELECT 1")
+                except Exception:
+                    # Connection dead, remove and retry
+                    self._connections.discard(conn)
+                    await conn.close()
+                    self._stats.total_errors += 1
+                    continue
+
+                self._stats.active_connections += 1
+                self._stats.total_acquired += 1
+                return conn
+
+            except PoolExhaustedError:
+                raise
+            except Exception as e:
+                self._stats.total_errors += 1
+                if attempt < self.max_retries - 1:
+                    log.warning(
+                        "Connection attempt %d failed: %s, retrying...",
+                        attempt + 1,
+                        e,
+                    )
+                    await asyncio.sleep(self.retry_delay)
+                else:
+                    raise DatabaseConnectionError(
+                        f"Failed to acquire connection after {self.max_retries} attempts"
+                    ) from e
+
+        raise DatabaseConnectionError("Failed to acquire connection")
+
+    async def release(self, conn: aiosqlite.Connection) -> None:
+        """Release a connection back to the pool.
+
+        Args:
+            conn: Connection to release.
+        """
+        if self._closed:
+            await conn.close()
+            return
+
+        self._stats.active_connections -= 1
+        self._stats.total_released += 1
+
+        # Return to pool
+        await self._pool.put(conn)
+        self._stats.idle_connections += 1
+
+    @asynccontextmanager
+    async def connection(self) -> AsyncGenerator[aiosqlite.Connection, None]:
+        """Context manager for acquiring and releasing connections.
+
+        Yields:
+            Database connection.
+
+        Examples:
+            >>> async with pool.connection() as conn:  # doctest: +SKIP
+            ...     await conn.execute("SELECT 1")
+        """
+        conn = await self.acquire()
+        try:
+            yield conn
+        finally:
+            await self.release(conn)
+
+    async def close(self) -> None:
+        """Close all connections and shutdown the pool."""
+        self._closed = True
+
+        async with self._lock:
+            # Close all connections
+            for conn in self._connections:
+                try:
+                    await conn.close()
+                except Exception:
+                    # Connection may be already closed - intentional silent catch
+                    log.debug("Failed to close connection (may be already closed)", exc_info=True)
+            self._connections.clear()
+
+            # Empty the queue
+            while not self._pool.empty():
+                try:
+                    self._pool.get_nowait()
+                except asyncio.QueueEmpty:
+                    break
+
+        self._stats.active_connections = 0
+        self._stats.idle_connections = 0
+        log.debug("Connection pool closed")
+
+    @property
+    def stats(self) -> PoolStats:
+        """Get pool statistics."""
+        return self._stats
+
+    @property
+    def size(self) -> int:
+        """Current number of connections in pool."""
+        return len(self._connections)
+
+    @property
+    def is_closed(self) -> bool:
+        """Whether the pool is closed."""
+        return self._closed
+
+
+__all__ = ["ConnectionPool", "PoolStats"]