PWF/CoreModules/database.py

"""SQLite数据库操作模块 - 使用标准库sqlite3"""
import sqlite3
import time
from typing import Any, Optional
from ..Convention.Runtime.Config import *
from ..Convention.Runtime.GlobalConfig import ProjectConfig, ConsoleFrontColor
from ..Convention.Runtime.Architecture import Architecture
from ..Convention.Runtime.File import ToolFile

logger: ProjectConfig = Architecture.Get(ProjectConfig)
DATABASE_PATH = logger.GetFile(logger.FindItem("database_path", "db.db"), False).GetFullPath()
logger.SaveProperties()

SCHEDULED_TASK_TABLE = "scheduled_tasks"
STATUS_PENDING = "pending"
STATUS_RUNNING = "running"
STATUS_COMPLETED = "completed"
STATUS_FAILED = "failed"

class Database:
    """数据库管理类"""
    
    def __init__(self, db_path: str = DATABASE_PATH):
        """初始化数据库连接
        
        Args:
            db_path: 数据库文件路径
        """
        self.db_path = db_path
        self._conn: Optional[sqlite3.Connection] = None
        self._ensure_db_exists()
        self.init_tables()
        Architecture.Register(Database, self, lambda: None)
    
    def _ensure_db_exists(self):
        """确保数据库目录存在"""
        db_dir = ToolFile(self.db_path).BackToParentDir()
        db_dir.MustExistsPath()
    
    @property
    def conn(self) -> sqlite3.Connection:
        """获取数据库连接（懒加载）"""
        if self._conn is None:
            try:
                self._conn = sqlite3.connect(
                    self.db_path,
                    check_same_thread=False,  # 允许多线程访问
                    isolation_level=None,  # 自动提交
                    timeout=30.0  # 增加超时时间
                )
                self._conn.row_factory = sqlite3.Row  # 支持字典式访问
                
                # 启用WAL模式以提高并发性能
                self._conn.execute("PRAGMA journal_mode=WAL")
                self._conn.execute("PRAGMA synchronous=NORMAL")
                self._conn.execute("PRAGMA cache_size=1000")
                self._conn.execute("PRAGMA temp_store=MEMORY")
                
                logger.Log("Info", f"{ConsoleFrontColor.GREEN}数据库连接成功: {self.db_path}{ConsoleFrontColor.RESET}")
            except Exception as e:
                logger.Log("Error", f"{ConsoleFrontColor.RED}数据库连接失败: {e}{ConsoleFrontColor.RESET}")
                raise
        return self._conn
    
    def _table_exists(self, table_name: str) -> bool:
        """检查表是否存在
        
        Args:
            table_name: 表名
            
        Returns:
            是否存在
        """
        cursor = self.conn.cursor()
        cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name=?", (table_name,))
        return cursor.fetchone() is not None

    def define_table(self, table_name: str):
        """定义表
        
        Args:
            table_name: 表名
        """
        if not self._table_exists(table_name):
            cursor = self.conn.cursor()
            cursor.execute(f"CREATE TABLE IF NOT EXISTS {table_name} (id INTEGER PRIMARY KEY AUTOINCREMENT)")
            logger.Log("Info", f"{ConsoleFrontColor.GREEN}为表 {table_name} 创建{ConsoleFrontColor.RESET}")
        return self

    def _column_exists(self, table_name: str, column_name: str) -> bool:
        """检查表中列是否存在
        
        Args:
            table_name: 表名
            column_name: 列名
            
        Returns:
            是否存在
        """
        cursor = self.conn.cursor()
        cursor.execute(f"PRAGMA table_info({table_name})")
        columns = [row[1] for row in cursor.fetchall()]
        return column_name in columns
    
    def _add_column_if_not_exists(self, table_name: str, column_name: str, column_def: str):
        """安全地添加列（如果不存在）
        
        Args:
            table_name: 表名
            column_name: 列名
            column_def: 列定义（如 "INTEGER" 或 "TEXT DEFAULT ''"）
        """
        if not self._column_exists(table_name, column_name):
            try:
                cursor = self.conn.cursor()
                cursor.execute(f"ALTER TABLE {table_name} ADD COLUMN {column_name} {column_def}")
                logger.Log("Info", f"{ConsoleFrontColor.GREEN}为表 {table_name} 添加列 {column_name}{ConsoleFrontColor.RESET}")
            except Exception as e:
                logger.Log("Warning", f"{ConsoleFrontColor.YELLOW}添加列失败: {e}{ConsoleFrontColor.RESET}")
    
    def define_column(self, table_name: str, column_name: str, column_def: str):
        """定义列
        
        Args:
            table_name: 表名
            column_name: 列名
            column_def: 列定义（如 "INTEGER" 或 "TEXT DEFAULT ''"）
        """
        self._add_column_if_not_exists(table_name, column_name, column_def)
        return self

    def init_tables(self):
        """初始化数据库表"""
        self._ensure_scheduled_tasks_table()
    
    def _ensure_scheduled_tasks_table(self) -> None:
        cursor = self.conn.cursor()
        cursor.execute(
            f"""
            CREATE TABLE IF NOT EXISTS {SCHEDULED_TASK_TABLE} (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                plugin_module TEXT NOT NULL,
                plugin_class TEXT,
                callback_name TEXT NOT NULL,
                payload TEXT,
                execute_at INTEGER NOT NULL,
                status TEXT NOT NULL DEFAULT '{STATUS_PENDING}',
                attempts INTEGER NOT NULL DEFAULT 0,
                last_error TEXT,
                created_at INTEGER NOT NULL,
                updated_at INTEGER NOT NULL
            )
            """
        )
        cursor.execute(
            f"CREATE INDEX IF NOT EXISTS idx_{SCHEDULED_TASK_TABLE}_status_execute_at ON {SCHEDULED_TASK_TABLE}(status, execute_at)"
        )

    def create_scheduled_task(
        self,
        plugin_module: str,
        plugin_class: Optional[str],
        callback_name: str,
        payload: Optional[str],
        execute_at_ms: int
    ) -> int:
        now_ms = int(time.time() * 1000)
        cursor = self.conn.cursor()
        cursor.execute(
            f"""
            INSERT INTO {SCHEDULED_TASK_TABLE} (
                plugin_module, plugin_class, callback_name, payload,
                execute_at, status, attempts, created_at, updated_at
            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
            """,
            (
                plugin_module,
                plugin_class,
                callback_name,
                payload,
                execute_at_ms,
                STATUS_PENDING,
                0,
                now_ms,
                now_ms,
            )
        )
        task_id = cursor.lastrowid
        return int(task_id)

    def get_due_tasks(self, now_ms: int, limit: int) -> list[sqlite3.Row]:
        cursor = self.conn.cursor()
        cursor.execute(
            f"""
            SELECT * FROM {SCHEDULED_TASK_TABLE}
            WHERE status = ? AND execute_at <= ?
            ORDER BY execute_at ASC
            LIMIT ?
            """,
            (STATUS_PENDING, now_ms, limit)
        )
        return cursor.fetchall()

    def update_task_status(
        self,
        task_id: int,
        status: str,
        *,
        attempts: Optional[int] = None,
        last_error: Optional[str] = None
    ) -> None:
        now_ms = int(time.time() * 1000)
        sets = ["status = ?"]
        params: list[Any] = [status]
        if attempts is not None:
            sets.append("attempts = ?")
            params.append(attempts)
        if last_error is not None:
            sets.append("last_error = ?")
            params.append(last_error)
        params.append(task_id)
        set_clause = ", ".join(sets)
        cursor = self.conn.cursor()
        cursor.execute(
            f"UPDATE {SCHEDULED_TASK_TABLE} SET {set_clause} WHERE id = ?",
            params
        )

    def reset_running_tasks(self) -> None:
        cursor = self.conn.cursor()
        now_ms = int(time.time() * 1000)
        cursor.execute(
            f"""
            UPDATE {SCHEDULED_TASK_TABLE}
            SET status = ?, updated_at = ?
            WHERE status = ?
            """,
            (STATUS_PENDING, now_ms, STATUS_RUNNING)
        )
    
    def close(self):
        """关闭数据库连接"""
        if self._conn:
            self._conn.close()
            self._conn = None
            logger.Log("Info", f"{ConsoleFrontColor.GREEN}数据库连接已关闭{ConsoleFrontColor.RESET}")


def get_db() -> Database:
    """获取全局数据库实例（单例模式）"""
    if not Architecture.Contains(Database):
        return Database()
    return Architecture.Get(Database)

__all__ = [
    "get_db",
    "Database",
    "SCHEDULED_TASK_TABLE",
    "STATUS_PENDING",
    "STATUS_RUNNING",
    "STATUS_COMPLETED",
    "STATUS_FAILED",
]