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 VIRTUAL TABLE IF NOT EXISTS wiki_fts USING fts5(
                path UNINDEXED,
                title,
                content,
                summary,
                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
        ))

        # 更新标签关联
        await self.update_page_tags(page.path, page.tags)

    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 搜索 ===

    def _validate_fts_query(self, query: str) -> str:
        """
        验证和清理 FTS5 查询字符串，防止注入攻击

        FTS5 支持的特殊字符：
        - 双引号：短语查询
        - *：前缀查询
        - AND, OR, NOT：布尔运算符

        验证规则：
        - 移除不安全的控制字符
        - 限制查询长度（防止 DoS）
        - 转义双引号防止短语注入
        """
        if not query:
            raise ValueError("Query cannot be empty")

        # 1. 限制查询长度
        max_query_length = 500
        if len(query) > max_query_length:
            logger.warning(f"Query too long ({len(query)} chars), truncating to {max_query_length}")
            query = query[:max_query_length]

        # 2. 移除控制字符（除了换行、制表符）
        import re
        query = re.sub(r'[\x00-\x08\x0b-\x0c\x0e-\x1f\x7f]', '', query)

        # 3. 转义未闭合的双引号（防止短语查询注入）
        # 计算引号数量，如果是奇数则转义最后一个
        quote_count = query.count('"')
        if quote_count % 2 != 0:
            # 找到最后一个引号并转义
            last_quote_idx = query.rfind('"')
            query = query[:last_quote_idx] + '\\"' + query[last_quote_idx+1:]

        # 4. 防止布尔运算符注入（移除前后空格的运算符）
        # 这是为了防止类似 "term AND DROP TABLE" 的攻击
        # FTS5 会在查询语法错误时返回空结果，但我们需要额外保护
        query = re.sub(r'\s+(AND|OR|NOT)\s+', ' ', query, flags=re.IGNORECASE)

        return query.strip()

    async def fts_search(self, query: str, limit: int = 10, with_snippets: bool = False) -> List[WikiPage]:
        """FTS5 全文搜索

        Args:
            query: 搜索查询
            limit: 返回结果数量限制
            with_snippets: 是否返回匹配片段（暂未实现，预留接口）
        """
        # 验证和清理查询
        try:
            safe_query = self._validate_fts_query(query)
        except ValueError as e:
            logger.warning(f"Invalid FTS query: {e}")
            return []

        # 先从 FTS5 获取匹配的路径
        try:
            fts_rows = await self.fetch_all(
                "SELECT path FROM wiki_fts WHERE wiki_fts MATCH ? ORDER BY rank LIMIT ?",
                (safe_query, limit)
            )
        except aiosqlite.OperationalError as e:
            # FTS5 语法错误时返回空结果（不应该崩溃）
            logger.warning(f"FTS5 query failed: {e}")
            return []

        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 update_page_tags(self, path: str, tags: List[str]) -> None:
        """更新页面的标签关联

        先删除旧的标签关联，再插入新的。
        同时更新 wiki_tags 表中的计数。
        """
        # 删除旧的标签关联
        await self.execute_with_retry("DELETE FROM wiki_page_tags WHERE path = ?", (path,))

        # 插入新的标签关联
        for tag in tags:
            await self.execute_with_retry(
                "INSERT INTO wiki_page_tags (path, tag) VALUES (?, ?)",
                (path, tag)
            )
            # 更新标签计数
            await self.upsert_tag(tag)

    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 内容

        FTS5 虚拟表不支持 UPSERT，使用 DELETE + INSERT 替代。
        所有内容（包括 content 字段）直接存储在 FTS 表中。

        优化：连字符扩展（multi-agent → multi agent）以提升搜索召回率
        """
        # 先删除旧记录（如果存在）
        await self.execute_with_retry("DELETE FROM wiki_fts WHERE path = ?", (path,))

        # 扩展连字符词汇（multi-agent → multi agent）
        # 这样搜索 "multi" 或 "agent" 都能匹配到
        expanded_title = self._expand_hyphenated_text(title)
        expanded_summary = self._expand_hyphenated_text(summary)

        # 插入扩展后的 title 和 summary（原版保留在 content 中）
        await self.execute_with_retry(
            "INSERT INTO wiki_fts (path, title, summary, content) VALUES (?, ?, ?, ?)",
            (path, f"{title} {expanded_title}", f"{summary} {expanded_summary}", content)
        )

    def _expand_hyphenated_text(self, text: str) -> str:
        """扩展连字符词汇

        multi-agent system → multi agent system
        这样搜索 "multi" 或 "agent" 都能匹配

        Args:
            text: 原始文本

        Returns:
            扩展后的文本（连字符替换为空格，并去除首尾空白）
        """
        if not text:
            return ""
        return text.replace('-', ' ').strip()

    async def delete_fts_content(self, path: str) -> None:
        """删除 FTS5 内容"""
        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()