fastapi_app/lib/core/database.py

"""
Database manager for SQLite file metadata system.

Provides connection management, transactions, and thread-safe database access.
Uses context managers for proper resource cleanup.
"""

import sqlite3
import queue
from contextlib import contextmanager
from pathlib import Path
from typing import Optional, Generator
from fastapi_app.lib.core.db_schema import initialize_database
from . import sqlite_utils


class DatabaseManager:
    """
    Manages SQLite database connections and transactions.

    Thread-safe connection management with context managers for
    automatic resource cleanup and transaction handling.
    """

    def __init__(self, db_path: Path, logger=None):
        """
        Initialize database manager.

        Args:
            db_path: Path to SQLite database file
            logger: Optional logger instance
        """
        self.db_path = db_path
        self.logger = logger
        self._ensure_db_exists()
        self._pool = queue.Queue()

    def _ensure_db_exists(self) -> None:
        """
        Ensure database file and schema exist.

        Creates database file and initializes schema if needed.
        Runs any pending migrations automatically.
        This method is idempotent - safe to call multiple times.

        Uses a per-database lock to prevent concurrent schema initialization
        which can corrupt the database.
        """
        # Ensure parent directory exists
        self.db_path.parent.mkdir(parents=True, exist_ok=True)

        # Use per-database lock to prevent concurrent schema initialization
        with sqlite_utils.with_db_lock(self.db_path):
            # Use raw connection to avoid recursive locking from sqlite_utils.get_connection
            # and to ensure we have control over WAL mode setting
            conn = sqlite3.connect(str(self.db_path), timeout=60.0, isolation_level=None)
            try:
                # Enable WAL mode explicitly
                conn.execute("PRAGMA journal_mode = WAL")
                conn.execute("PRAGMA foreign_keys = ON")

                # Create database and initialize schema (including migrations)
                initialize_database(conn, self.logger, db_path=self.db_path)
            finally:
                conn.close()

    @contextmanager
    def get_connection(self) -> Generator[sqlite3.Connection, None, None]:
        """
        Context manager for database connections.

        Yields a connection with row_factory set to sqlite3.Row
        for dict-like access to query results.

        Usage:
            with db_manager.get_connection() as conn:
                cursor = conn.cursor()
                cursor.execute("SELECT * FROM files")

        Yields:
            sqlite3.Connection: Database connection
        """
        try:
            conn = self._pool.get(block=False)
        except queue.Empty:
            conn = sqlite3.connect(str(self.db_path), timeout=60.0, check_same_thread=False, isolation_level=None)
            conn.row_factory = sqlite3.Row
            conn.execute("PRAGMA foreign_keys = ON")

        try:
            yield conn
        finally:
            # Rollback any uncommitted changes to ensure clean state for next use
            try:
                conn.rollback()
            except sqlite3.OperationalError:
                pass
            self._pool.put(conn)

    @contextmanager
    def transaction(self) -> Generator[sqlite3.Connection, None, None]:
        """
        Context manager for database transactions.

        Automatically commits on success, rolls back on exception.
        Useful for operations that require atomicity.

        Usage:
            with db_manager.transaction() as conn:
                cursor = conn.cursor()
                cursor.execute("INSERT INTO files ...")
                cursor.execute("UPDATE files ...")
                # Auto-commit on exit (or rollback on exception)

        Yields:
            sqlite3.Connection: Database connection with transaction
        """
        with self.get_connection() as conn:
            conn.execute("BEGIN")
            try:
                yield conn
                conn.commit()
                if self.logger:
                    self.logger.debug("Transaction committed")
            except Exception:
                try:
                    conn.rollback()
                except sqlite3.OperationalError:
                    pass
                raise

    def execute_query(
        self,
        query: str,
        params: tuple = (),
        fetch_one: bool = False
    ) -> Optional[list | dict]:
        """
        Execute a SELECT query and return results.

        Convenience method for simple queries.

        Args:
            query: SQL query string
            params: Query parameters
            fetch_one: If True, return single row; if False, return all rows

        Returns:
            Single row (dict) if fetch_one=True, list of rows otherwise
        """
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute(query, params)

            if fetch_one:
                row = cursor.fetchone()
                return dict(row) if row else None
            else:
                rows = cursor.fetchall()
                return [dict(row) for row in rows]

    def execute_update(
        self,
        query: str,
        params: tuple = ()
    ) -> int:
        """
        Execute an INSERT, UPDATE, or DELETE query.

        Convenience method for simple updates.

        Args:
            query: SQL query string
            params: Query parameters

        Returns:
            Number of affected rows
        """
        with self.transaction() as conn:
            cursor = conn.cursor()
            cursor.execute(query, params)
            return cursor.rowcount

    def clear_all_data(self) -> None:
        """
        Clear all data from the files table.

        This removes all file metadata records from the database.
        The table schema remains intact.

        Warning: This operation cannot be undone.
        """
        with self.transaction() as conn:
            cursor = conn.cursor()
            cursor.execute("DELETE FROM files")
            deleted_count = cursor.rowcount

            if self.logger:
                self.logger.info(f"Cleared {deleted_count} records from database")