sanguo_llmwiki/mcp_server/storage/database.py

"""
Storage Layer - Database 类

使用 aiosqlite 实现真正的异步支持，配合 WAL 模式和并发保护。

参考设计文档：第 2.4 节
"""

import aiosqlite
import asyncio
import logging
import time
from typing import List, Tuple, Dict, Set, Optional, Any
from datetime import datetime
from dataclasses import dataclass, asdict
import hashlib
import json

logger = logging.getLogger(__name__)


@dataclass
class WikiPage:
    """Wiki 页面模型"""
    path: str
    title: str
    category: str
    tags: List[str]
    summary: str
    content_hash: str
    lifecycle: str
    source_tool: str
    created_at: str
    updated_at: str
    indexed_at: str

    def is_stale(self, days: int = 90) -> bool:
        """检查页面是否过期"""
        updated = datetime.fromisoformat(self.updated_at)
        return (datetime.now() - updated).days > days

    def to_dict(self) -> dict:
        """转换为字典"""
        return asdict(self)


class Database:
    """SQLite 数据库封装，使用 aiosqlite 实现真正的异步支持"""

    def __init__(self, path: str):
        self.path = path
        self._conn: Optional[aiosqlite.Connection] = None
        self._lock = asyncio.Lock()

    async def connect(self) -> None:
        """建立连接，启用 WAL 模式和并发保护"""
        logger.info(f"Connecting to database: {self.path}")
        self._conn = await aiosqlite.connect(self.path)

        # 启用 WAL 模式提升并发性能
        await self._conn.execute("PRAGMA journal_mode=WAL")
        await self._conn.execute("PRAGMA busy_timeout=10000")  # 10s
        await self._conn.execute("PRAGMA synchronous=NORMAL")
        await self._conn.commit()

        # 创建表结构
        await self._create_tables()

        logger.info("Database connected and initialized")

    async def _create_tables(self) -> None:
        """创建所有表结构"""

        # 页面索引表
        await self._conn.execute("""
            CREATE TABLE IF NOT EXISTS wiki_pages (
                path TEXT PRIMARY KEY,
                title TEXT NOT NULL,
                category TEXT,
                tags TEXT,
                summary TEXT,
                content_hash TEXT NOT NULL,
                lifecycle TEXT DEFAULT 'draft',
                source_tool TEXT DEFAULT 'other',
                created_at TIMESTAMP,
                updated_at TIMESTAMP,
                indexed_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        """)

        # 索引
        await self._conn.execute("CREATE INDEX IF NOT EXISTS idx_pages_category ON wiki_pages(category)")
        await self._conn.execute("CREATE INDEX IF NOT EXISTS idx_pages_lifecycle ON wiki_pages(lifecycle)")
        await self._conn.execute("CREATE INDEX IF NOT EXISTS idx_pages_updated ON wiki_pages(updated_at)")
        await self._conn.execute("CREATE INDEX IF NOT EXISTS idx_pages_source_tool ON wiki_pages(source_tool)")

        # 内容表（FTS5 外部内容表）- 必须先创建
        await self._conn.execute("""
            CREATE TABLE IF NOT EXISTS wiki_content (
                path TEXT PRIMARY KEY,
                content TEXT NOT NULL
            )
        """)

        # FTS5 全文搜索表
        await self._conn.execute("""
            CREATE VIRTUAL TABLE IF NOT EXISTS wiki_fts USING fts5(
                path UNINDEXED,
                title,
                content,
                summary,
                content=wiki_content,
                content_rowid=rowid,
                tokenize = 'porter unicode61'
            )
        """)

        # 链接关系表
        await self._conn.execute("""
            CREATE TABLE IF NOT EXISTS wiki_links (
                source TEXT NOT NULL,
                target TEXT NOT NULL,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                PRIMARY KEY (source, target)
            )
        """)
        await self._conn.execute("CREATE INDEX IF NOT EXISTS idx_links_target ON wiki_links(target)")

        # 标签索引表
        await self._conn.execute("""
            CREATE TABLE IF NOT EXISTS wiki_tags (
                tag TEXT PRIMARY KEY,
                count INTEGER DEFAULT 0,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        """)

        # 页面-标签关联表
        await self._conn.execute("""
            CREATE TABLE IF NOT EXISTS wiki_page_tags (
                path TEXT NOT NULL,
                tag TEXT NOT NULL,
                PRIMARY KEY (path, tag),
                FOREIGN KEY (path) REFERENCES wiki_pages(path) ON DELETE CASCADE,
                FOREIGN KEY (tag) REFERENCES wiki_tags(tag) ON DELETE CASCADE
            )
        """)
        await self._conn.execute("CREATE INDEX IF NOT EXISTS idx_page_tags_tag ON wiki_page_tags(tag)")

        # 索引元数据表
        await self._conn.execute("""
            CREATE TABLE IF NOT EXISTS wiki_meta (
                key TEXT PRIMARY KEY,
                value TEXT,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        """)

        await self._conn.commit()

    async def execute(self, sql: str, params: tuple = ()) -> aiosqlite.Cursor:
        """执行 SQL（带写入锁）"""
        async with self._lock:
            cursor = await self._conn.execute(sql, params)
            await self._conn.commit()
            return cursor

    async def execute_with_retry(self, sql: str, params: tuple = (), max_retries: int = 2) -> aiosqlite.Cursor:
        """带重试的数据库操作（指数退避）"""
        for attempt in range(max_retries):
            try:
                return await self.execute(sql, params)
            except aiosqlite.OperationalError as e:
                if "database is locked" in str(e) and attempt < max_retries - 1:
                    wait_time = 0.1 * (2 ** attempt)
                    logger.warning(f"Database locked, retrying in {wait_time}s (attempt {attempt + 1}/{max_retries})")
                    await asyncio.sleep(wait_time)
                    continue
                raise

    async def fetch_all(self, sql: str, params: tuple = ()) -> List[Tuple]:
        """查询所有结果（读操作无需锁，WAL 自动处理）"""
        cursor = await self._conn.execute(sql, params)
        return await cursor.fetchall()

    async def fetch_one(self, sql: str, params: tuple = ()) -> Optional[Tuple]:
        """查询单个结果"""
        rows = await self.fetch_all(sql, params)
        return rows[0] if rows else None

    # === 页面操作 ===

    async def get_page(self, path: str) -> Optional[WikiPage]:
        """获取单个页面"""
        row = await self.fetch_one(
            "SELECT path, title, category, tags, summary, content_hash, lifecycle, source_tool, created_at, updated_at, indexed_at FROM wiki_pages WHERE path = ?",
            (path,)
        )
        if not row:
            return None

        return WikiPage(
            path=row[0],
            title=row[1],
            category=row[2],
            tags=json.loads(row[3]) if row[3] else [],
            summary=row[4],
            content_hash=row[5],
            lifecycle=row[6],
            source_tool=row[7],
            created_at=row[8],
            updated_at=row[9],
            indexed_at=row[10]
        )

    async def get_all_pages(self) -> List[WikiPage]:
        """获取所有页面"""
        rows = await self.fetch_all(
            "SELECT path, title, category, tags, summary, content_hash, lifecycle, source_tool, created_at, updated_at, indexed_at FROM wiki_pages"
        )
        return [
            WikiPage(
                path=row[0],
                title=row[1],
                category=row[2],
                tags=json.loads(row[3]) if row[3] else [],
                summary=row[4],
                content_hash=row[5],
                lifecycle=row[6],
                source_tool=row[7],
                created_at=row[8],
                updated_at=row[9],
                indexed_at=row[10]
            )
            for row in rows
        ]

    async def get_page_hash(self, path: str) -> Optional[str]:
        """获取页面内容哈希"""
        row = await self.fetch_one("SELECT content_hash FROM wiki_pages WHERE path = ?", (path,))
        return row[0] if row else None

    async def upsert_page(self, page: WikiPage) -> None:
        """插入或更新页面"""
        await self.execute_with_retry("""
            INSERT INTO wiki_pages (path, title, category, tags, summary, content_hash, lifecycle, source_tool, created_at, updated_at, indexed_at)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            ON CONFLICT(path) DO UPDATE SET
                title = excluded.title,
                category = excluded.category,
                tags = excluded.tags,
                summary = excluded.summary,
                content_hash = excluded.content_hash,
                lifecycle = excluded.lifecycle,
                source_tool = excluded.source_tool,
                updated_at = excluded.updated_at,
                indexed_at = excluded.indexed_at
        """, (
            page.path, page.title, page.category, json.dumps(page.tags),
            page.summary, page.content_hash, page.lifecycle, page.source_tool,
            page.created_at, page.updated_at, page.indexed_at
        ))

    async def delete_page(self, path: str) -> None:
        """删除页面（级联删除相关数据）"""
        await self.execute_with_retry("DELETE FROM wiki_pages WHERE path = ?", (path,))

    async def get_all_indexed_paths(self) -> Set[str]:
        """获取所有已索引的页面路径"""
        rows = await self.fetch_all("SELECT path FROM wiki_pages")
        return {row[0] for row in rows}

    # === FTS5 搜索 ===

    async def fts_search(self, query: str, limit: int = 10) -> List[WikiPage]:
        """FTS5 全文搜索"""
        # 先从 FTS5 获取匹配的路径
        fts_rows = await self.fetch_all(
            "SELECT path FROM wiki_fts WHERE wiki_fts MATCH ? ORDER BY rank LIMIT ?",
            (query, limit)
        )
        paths = [row[0] for row in fts_rows]

        if not paths:
            return []

        # 从 wiki_pages 获取完整数据
        placeholders = ','.join('?' * len(paths))
        rows = await self.fetch_all(
            f"SELECT path, title, category, tags, summary, content_hash, lifecycle, source_tool, created_at, updated_at, indexed_at FROM wiki_pages WHERE path IN ({placeholders})",
            paths
        )

        return [
            WikiPage(
                path=row[0],
                title=row[1],
                category=row[2],
                tags=json.loads(row[3]) if row[3] else [],
                summary=row[4],
                content_hash=row[5],
                lifecycle=row[6],
                source_tool=row[7],
                created_at=row[8],
                updated_at=row[9],
                indexed_at=row[10]
            )
            for row in rows
        ]

    async def search_by_tags(self, tags: List[str]) -> List[WikiPage]:
        """按标签搜索"""
        if not tags:
            return []

        placeholders = ','.join('?' * len(tags))
        rows = await self.fetch_all(
            f"SELECT DISTINCT p.path, p.title, p.category, p.tags, p.summary, p.content_hash, p.lifecycle, p.source_tool, p.created_at, p.updated_at, p.indexed_at FROM wiki_pages p JOIN wiki_page_tags pt ON p.path = pt.path WHERE pt.tag IN ({placeholders})",
            tags
        )

        return [
            WikiPage(
                path=row[0],
                title=row[1],
                category=row[2],
                tags=json.loads(row[3]) if row[3] else [],
                summary=row[4],
                content_hash=row[5],
                lifecycle=row[6],
                source_tool=row[7],
                created_at=row[8],
                updated_at=row[9],
                indexed_at=row[10]
            )
            for row in rows
        ]

    async def search_by_source_tool(self, tool_name: str, limit: int = 50) -> List[WikiPage]:
        """按来源工具搜索（memory_bridge 使用）"""
        rows = await self.fetch_all(
            "SELECT path, title, category, tags, summary, content_hash, lifecycle, source_tool, created_at, updated_at, indexed_at FROM wiki_pages WHERE source_tool = ? ORDER BY updated_at DESC LIMIT ?",
            (tool_name, limit)
        )

        return [
            WikiPage(
                path=row[0],
                title=row[1],
                category=row[2],
                tags=json.loads(row[3]) if row[3] else [],
                summary=row[4],
                content_hash=row[5],
                lifecycle=row[6],
                source_tool=row[7],
                created_at=row[8],
                updated_at=row[9],
                indexed_at=row[10]
            )
            for row in rows
        ]

    # === 链接操作 ===

    async def get_links(self, path: str) -> Set[str]:
        """获取页面的出链"""
        rows = await self.fetch_all("SELECT target FROM wiki_links WHERE source = ?", (path,))
        return {row[0] for row in rows}

    async def get_backlinks(self, path: str) -> Set[str]:
        """获取页面的反向链接"""
        rows = await self.fetch_all("SELECT source FROM wiki_links WHERE target = ?", (path,))
        return {row[0] for row in rows}

    async def upsert_link(self, source: str, target: str) -> None:
        """插入或更新链接"""
        await self.execute_with_retry(
            "INSERT INTO wiki_links (source, target) VALUES (?, ?) ON CONFLICT(source, target) DO NOTHING",
            (source, target)
        )

    async def delete_links(self, path: str) -> None:
        """删除页面的所有链接"""
        await self.execute_with_retry("DELETE FROM wiki_links WHERE source = ?", (path,))

    # === 标签操作 ===

    async def upsert_tag(self, tag: str, count: int = 1) -> None:
        """插入或更新标签"""
        await self.execute_with_retry(
            "INSERT INTO wiki_tags (tag, count) VALUES (?, ?) ON CONFLICT(tag) DO UPDATE SET count = count + excluded.count",
            (tag, count)
        )

    async def get_all_tags(self) -> Dict[str, int]:
        """获取所有标签"""
        rows = await self.fetch_all("SELECT tag, count FROM wiki_tags")
        return {row[0]: row[1] for row in rows}

    # === 元数据操作 ===

    async def get_meta(self, key: str, default: str = "") -> str:
        """获取元数据"""
        row = await self.fetch_one("SELECT value FROM wiki_meta WHERE key = ?", (key,))
        return row[0] if row else default

    async def set_meta(self, key: str, value: str) -> None:
        """设置元数据"""
        await self.execute_with_retry(
            "INSERT INTO wiki_meta (key, value) VALUES (?, ?) ON CONFLICT(key) DO UPDATE SET value = excluded.value",
            (key, value)
        )

    # === FTS5 内容更新 ===

    async def update_fts_content(self, path: str, title: str, content: str, summary: str) -> None:
        """更新 FTS5 内容"""
        # 先更新内容表
        await self.execute_with_retry(
            "INSERT INTO wiki_content (path, content) VALUES (?, ?) ON CONFLICT(path) DO UPDATE SET content = excluded.content",
            (path, content)
        )

        # 再更新 FTS5 表（会自动从内容表同步）
        await self.execute_with_retry(
            "INSERT INTO wiki_fts (path, title, summary) VALUES (?, ?, ?) ON CONFLICT(path) DO UPDATE SET title = excluded.title, summary = excluded.summary",
            (path, title, summary)
        )

    async def delete_fts_content(self, path: str) -> None:
        """删除 FTS5 内容"""
        await self.execute_with_retry("DELETE FROM wiki_content WHERE path = ?", (path,))
        await self.execute_with_retry("DELETE FROM wiki_fts WHERE path = ?", (path,))

    # === 状态查询 ===

    async def get_stats(self) -> Dict[str, Any]:
        """获取索引统计信息"""
        page_count = await self.fetch_one("SELECT COUNT(*) FROM wiki_pages")
        link_count = await self.fetch_one("SELECT COUNT(*) FROM wiki_links")
        tag_count = await self.fetch_one("SELECT COUNT(*) FROM wiki_tags")
        last_indexed = await self.fetch_one("SELECT MAX(indexed_at) FROM wiki_pages")

        return {
            "total_pages": page_count[0] if page_count else 0,
            "total_links": link_count[0] if link_count else 0,
            "total_tags": tag_count[0] if tag_count else 0,
            "last_indexed": last_indexed[0] if last_indexed and last_indexed[0] else None
        }

    async def get_recent_pages(self, days: int = 7, limit: int = 50) -> List[WikiPage]:
        """获取最近更新的页面"""
        rows = await self.fetch_all(
            "SELECT path, title, category, tags, summary, content_hash, lifecycle, source_tool, created_at, updated_at, indexed_at FROM wiki_pages WHERE updated_at >= datetime('now', '-' || ? || ' days') ORDER BY updated_at DESC LIMIT ?",
            (str(days), limit)
        )

        return [
            WikiPage(
                path=row[0],
                title=row[1],
                category=row[2],
                tags=json.loads(row[3]) if row[3] else [],
                summary=row[4],
                content_hash=row[5],
                lifecycle=row[6],
                source_tool=row[7],
                created_at=row[8],
                updated_at=row[9],
                indexed_at=row[10]
            )
            for row in rows
        ]

    async def get_new_tags(self, days: int = 7) -> List[Tuple[str, str]]:
        """获取新增标签"""
        rows = await self.fetch_all(
            "SELECT tag, updated_at FROM wiki_tags WHERE updated_at >= datetime('now', '-' || ? || ' days') ORDER BY updated_at DESC",
            (str(days),)
        )
        return list(rows)

    # === 完整性检查 ===

    async def check_integrity(self) -> bool:
        """检查数据库完整性"""
        try:
            result = await self.fetch_one("PRAGMA integrity_check")
            if result and result[0] == "ok":
                return True
            logger.warning(f"Database integrity check failed: {result}")
            return False
        except Exception as e:
            logger.error(f"Integrity check error: {e}")
            return False

    # === 关闭连接 ===

    async def close(self) -> None:
        """关闭连接"""
        if self._conn:
            await self._conn.close()
            self._conn = None
            logger.info("Database connection closed")


async def fix_dirty_states(db: Database) -> None:
    """启动时清理可能的脏状态"""
    logger.info("Checking for dirty states...")

    try:
        # 1. 检查 WAL 文件是否损坏
        if db._conn:
            await db._conn.execute("PRAGMA wal_checkpoint(PASSIVE)")
            await db._conn.commit()

        # 2. 检查数据库完整性
        if not await db.check_integrity():
            raise Exception("Database integrity check failed")

        logger.info("No dirty states found, database is healthy")

    except Exception as e:
        logger.warning(f"Dirty states detected, attempting recovery: {e}")
        # TODO: 实现自动重建逻辑
        # await rebuild_index(db)


def compute_content_hash(content: str) -> str:
    """计算内容 MD5 哈希"""
    return hashlib.md5(content.encode('utf-8')).hexdigest()