aiagent/team_projects/b9b14925-f054-43d0-b2f0-6ed6b38dbe4b/Company-context-SQLiteFlutter-React-Nati/02-db_config.py

"""
家庭阅读激励工具 - 数据库连接管理
==================================
提供SQLite连接池、WAL模式配置、事务上下文管理器。
纯本地存储，零网络依赖。

使用方式:
    from db_config import get_db, transaction
    db = get_db()
    with transaction(db) as conn:
        conn.execute("INSERT INTO ...")

配置通过环境变量，所有值有合理默认值。
"""
import sqlite3
import os
import logging
from contextlib import contextmanager
from typing import Optional

logger = logging.getLogger(__name__)

# ============================================================
# 配置（环境变量优先，提供合理默认值）
# ============================================================
DB_DIR = os.environ.get("READING_APP_DB_DIR", os.path.dirname(os.path.abspath(__file__)))
DB_NAME = os.environ.get("READING_APP_DB_NAME", "reading_app.db")
DB_PATH = os.path.join(DB_DIR, DB_NAME)

# SQLite PRAGMA 配置
PRAGMAS: dict[str, str] = {
    "journal_mode": "WAL",          # WAL模式：读写并发、性能更好
    "foreign_keys": "ON",           # 强制外键约束
    "busy_timeout": "5000",         # 5秒忙等超时（避免SQLITE_BUSY）
    "cache_size": "-8000",          # 8MB页缓存（负数为KB）
    "synchronous": "NORMAL",        # WAL下NORMAL足够安全
    "temp_store": "MEMORY",         # 临时表存内存
    "mmap_size": "268435456",       # 256MB内存映射（提升大表读取）
}

# ============================================================
# 连接管理
# ============================================================

class DatabaseManager:
    """
    SQLite数据库单例管理器。
    线程安全：每个线程持有独立连接（SQLite check_same_thread=False时
    由调用方保证不跨线程共享连接）。
    """

    def __init__(self, db_path: str = DB_PATH):
        self.db_path = db_path
        self._ensure_db_dir()

    def _ensure_db_dir(self) -> None:
        """确保数据库文件目录存在"""
        db_dir = os.path.dirname(self.db_path)
        if db_dir and not os.path.exists(db_dir):
            os.makedirs(db_dir, exist_ok=True)

    def get_connection(self) -> sqlite3.Connection:
        """
        获取新数据库连接，自动配置PRAGMA和row_factory。
        每次调用返回新连接，由调用方负责关闭。
        """
        conn = sqlite3.connect(
            self.db_path,
            detect_types=sqlite3.PARSE_DECLTYPES,
            check_same_thread=False,
        )
        conn.row_factory = sqlite3.Row  # 支持列名访问
        conn.execute("PRAGMA journal_mode=WAL")
        conn.execute("PRAGMA foreign_keys=ON")
        conn.execute("PRAGMA busy_timeout=5000")
        return conn

    def initialize_schema(self) -> None:
        """
        初始化数据库Schema（幂等操作）。
        读取01-schema.sql并执行，CREATE TABLE IF NOT EXISTS保证幂等。
        """
        schema_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "01-schema.sql")
        if not os.path.exists(schema_path):
            raise FileNotFoundError(f"Schema文件不存在: {schema_path}")

        with open(schema_path, "r", encoding="utf-8") as f:
            schema_sql = f.read()

        conn = self.get_connection()
        try:
            conn.executescript(schema_sql)
            conn.commit()
            logger.info("数据库Schema初始化完成: %s", self.db_path)
        finally:
            conn.close()


# 全局数据库管理器实例
_db_manager: Optional[DatabaseManager] = None


def get_db_manager() -> DatabaseManager:
    """获取全局DatabaseManager单例"""
    global _db_manager
    if _db_manager is None:
        _db_manager = DatabaseManager()
    return _db_manager


def get_db() -> sqlite3.Connection:
    """获取数据库连接（快捷方法）"""
    return get_db_manager().get_connection()


def init_db() -> None:
    """初始化数据库（应用启动时调用一次）"""
    get_db_manager().initialize_schema()


# ============================================================
# 事务上下文管理器
# ============================================================

@contextmanager
def transaction(db: sqlite3.Connection):
    """
    事务上下文管理器，自动commit/rollback。

    用法:
        db = get_db()
        try:
            with transaction(db) as conn:
                conn.execute("INSERT INTO children (name) VALUES (?)", ("小明",))
                conn.execute("INSERT INTO points (...) VALUES (...)", (...))
            # 退出with时自动commit
        except Exception:
            # 自动rollback
            raise
        finally:
            db.close()
    """
    try:
        yield db
        db.commit()
    except Exception:
        db.rollback()
        raise


@contextmanager
def db_session() -> sqlite3.Connection:
    """
    连接+事务一体化的便捷上下文管理器。
    自动获取连接、开启事务、提交/回滚、关闭连接。

    用法:
        with db_session() as conn:
            row = conn.execute("SELECT * FROM children WHERE id=?", (1,)).fetchone()
    """
    db = get_db()
    try:
        yield db
        db.commit()
    except Exception:
        db.rollback()
        raise
    finally:
        db.close()


# ============================================================
# 辅助工具
# ============================================================

def row_to_dict(row: sqlite3.Row) -> dict:
    """将sqlite3.Row转为普通字典（便于JSON序列化）"""
    if row is None:
        return None
    return dict(row)


def rows_to_list(rows: list[sqlite3.Row]) -> list[dict]:
    """将Row列表转为字典列表"""
    return [dict(r) for r in rows]