""" 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 搜索 === 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) -> List[WikiPage]: """FTS5 全文搜索""" # 验证和清理查询 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 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()