"""content_db.py — SQLite 内容数据库操作类""" import json import sqlite3 from pathlib import Path from typing import Optional class ContentDB: """SQLite 内容数据库操作类 每个文档对应一个独立的 SQLite 数据库文件,存储文档的所有 Blocks。 """ def __init__(self, db_path: str): self.db_path = Path(db_path) self.conn: Optional[sqlite3.Connection] = None def connect(self): """连接数据库""" self.conn = sqlite3.connect(str(self.db_path)) self.conn.row_factory = sqlite3.Row return self def close(self): """关闭连接""" if self.conn: self.conn.close() self.conn = None def __enter__(self): """上下文管理器入口""" return self.connect() def __exit__(self, exc_type, exc_val, exc_tb): """上下文管理器退出""" self.close() def create_tables(self): """创建 document_blocks 表""" self.conn.execute(""" CREATE TABLE IF NOT EXISTS document_blocks ( id TEXT PRIMARY KEY, block_order INTEGER NOT NULL, type TEXT NOT NULL, level INTEGER DEFAULT 0, "index" INTEGER DEFAULT 0, content TEXT NOT NULL, word_style TEXT DEFAULT '', style TEXT DEFAULT '{}', metadata TEXT DEFAULT '{}' ) """) # 创建索引 self.conn.execute( 'CREATE INDEX IF NOT EXISTS idx_block_order ON document_blocks(block_order)' ) self.conn.execute( 'CREATE INDEX IF NOT EXISTS idx_type ON document_blocks(type)' ) self.conn.execute( 'CREATE INDEX IF NOT EXISTS idx_level ON document_blocks(level)' ) self.conn.commit() def insert_blocks(self, blocks: list[dict]): """批量插入 blocks""" for block in blocks: content = block['content'] if isinstance(content, (dict, list)): content = json.dumps(content, ensure_ascii=False) self.conn.execute(""" INSERT INTO document_blocks (id, block_order, type, level, "index", content, word_style, style, metadata) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( block['id'], block['block_order'], block['type'], block.get('level', 0), block.get('index', 0), content, block.get('word_style', ''), json.dumps(block.get('style', {}), ensure_ascii=False), json.dumps(block.get('metadata', {}), ensure_ascii=False) )) self.conn.commit() def get_blocks(self, order_by: str = 'block_order') -> list[dict]: """查询所有 blocks""" cursor = self.conn.execute(f""" SELECT * FROM document_blocks ORDER BY {order_by} """) rows = cursor.fetchall() return [self._row_to_dict(row) for row in rows] def get_block_by_id(self, block_id: str) -> Optional[dict]: """按 ID 查询单个 block""" cursor = self.conn.execute(""" SELECT * FROM document_blocks WHERE id = ? """, (block_id,)) row = cursor.fetchone() return self._row_to_dict(row) if row else None def update_block(self, block_id: str, updates: dict): """更新单个 block""" set_clauses = [] params = [] if 'content' in updates: content = updates['content'] if isinstance(content, (dict, list)): content = json.dumps(content, ensure_ascii=False) set_clauses.append('content = ?') params.append(content) if 'style' in updates: set_clauses.append('style = ?') params.append(json.dumps(updates['style'], ensure_ascii=False)) if 'word_style' in updates: set_clauses.append('word_style = ?') params.append(updates['word_style']) if 'metadata' in updates: set_clauses.append('metadata = ?') params.append(json.dumps(updates['metadata'], ensure_ascii=False)) if not set_clauses: return params.append(block_id) sql = f"UPDATE document_blocks SET {', '.join(set_clauses)} WHERE id = ?" self.conn.execute(sql, params) self.conn.commit() def delete_block(self, block_id: str): """删除单个 block""" self.conn.execute('DELETE FROM document_blocks WHERE id = ?', (block_id,)) self.conn.commit() def search_blocks(self, query: str, block_type: Optional[str] = None) -> list[dict]: """搜索 blocks""" sql = "SELECT * FROM document_blocks WHERE content LIKE ?" params = [f'%{query}%'] if block_type: sql += " AND type = ?" params.append(block_type) sql += " ORDER BY block_order" cursor = self.conn.execute(sql, params) rows = cursor.fetchall() return [self._row_to_dict(row) for row in rows] def get_headings(self) -> list[dict]: """获取所有标题块""" cursor = self.conn.execute(""" SELECT * FROM document_blocks WHERE type = 'heading' ORDER BY block_order """) rows = cursor.fetchall() return [self._row_to_dict(row) for row in rows] def get_stats(self) -> dict: """获取统计信息""" cursor = self.conn.execute(""" SELECT type, COUNT(*) as count FROM document_blocks GROUP BY type """) stats = {row['type']: row['count'] for row in cursor.fetchall()} cursor = self.conn.execute("SELECT COUNT(*) as total FROM document_blocks") total = cursor.fetchone()['total'] return { 'total': total, 'by_type': stats } def calculate_index_for_insert( self, after_block_id: str, new_type: str, new_level: int = 0 ) -> int: """计算插入 block 的 index(稀疏排序) Args: after_block_id: 在哪个 block 后插入(如果为 None 则插入到末尾) new_type: 新 block 的类型 new_level: 新 block 的级别(标题有效) Returns: 计算得到的 index """ if after_block_id: # 1. 获取 after_block 的 block_order after_block = self.get_block_by_id(after_block_id) if not after_block: raise ValueError(f"Block not found: {after_block_id}") after_order = after_block['block_order'] # 2. 查找下一个同类型同级别的 block if new_type == 'heading': # 标题:按 level 查询 cursor = self.conn.execute(""" SELECT "index" FROM document_blocks WHERE type = ? AND level = ? AND block_order > ? ORDER BY block_order LIMIT 1 """, (new_type, new_level, after_order)) else: # 其他类型:只按 type 查询 cursor = self.conn.execute(""" SELECT "index" FROM document_blocks WHERE type = ? AND block_order > ? ORDER BY block_order LIMIT 1 """, (new_type, after_order)) next_row = cursor.fetchone() if next_row: next_index = next_row['index'] # 3. 查找前一个同类型同级别的 block if new_type == 'heading': cursor = self.conn.execute(""" SELECT "index" FROM document_blocks WHERE type = ? AND level = ? AND block_order <= ? ORDER BY block_order DESC LIMIT 1 """, (new_type, new_level, after_order)) else: cursor = self.conn.execute(""" SELECT "index" FROM document_blocks WHERE type = ? AND block_order <= ? ORDER BY block_order DESC LIMIT 1 """, (new_type, after_order)) prev_row = cursor.fetchone() prev_index = prev_row['index'] if prev_row else -100 # 4. 计算中间值 gap = next_index - prev_index if gap <= 1: # 间隙不足,触发局部重排 self._rebalance_indexes_between( new_type, new_level, prev_index, next_index ) # 重新查询 if new_type == 'heading': cursor = self.conn.execute(""" SELECT "index" FROM document_blocks WHERE type = ? AND level = ? AND block_order > ? ORDER BY block_order LIMIT 1 """, (new_type, new_level, after_order)) else: cursor = self.conn.execute(""" SELECT "index" FROM document_blocks WHERE type = ? AND block_order > ? ORDER BY block_order LIMIT 1 """, (new_type, after_order)) next_row = cursor.fetchone() next_index = next_row['index'] if next_row else prev_index + 200 new_index = (prev_index + next_index) // 2 else: # 没有下一个同类 block,追加到最后 if new_type == 'heading': cursor = self.conn.execute(""" SELECT MAX("index") as max_index FROM document_blocks WHERE type = ? AND level = ? """, (new_type, new_level)) else: cursor = self.conn.execute(""" SELECT MAX("index") as max_index FROM document_blocks WHERE type = ? """, (new_type,)) row = cursor.fetchone() max_index = row['max_index'] if row['max_index'] is not None else -100 new_index = max_index + 100 else: # 插入到文档末尾 if new_type == 'heading': cursor = self.conn.execute(""" SELECT MAX("index") as max_index FROM document_blocks WHERE type = ? AND level = ? """, (new_type, new_level)) else: cursor = self.conn.execute(""" SELECT MAX("index") as max_index FROM document_blocks WHERE type = ? """, (new_type,)) row = cursor.fetchone() max_index = row['max_index'] if row['max_index'] is not None else -100 new_index = max_index + 100 return new_index def calculate_block_order_for_insert(self, after_block_id: str = None) -> int: """计算插入 block 的 block_order(稀疏排序) Args: after_block_id: 在哪个 block 后插入(如果为 None 则插入到末尾) Returns: 计算得到的 block_order """ if after_block_id: # 1. 获取 after_block 的 block_order after_block = self.get_block_by_id(after_block_id) if not after_block: raise ValueError(f"Block not found: {after_block_id}") after_order = after_block['block_order'] # 2. 查询下一个 block 的 block_order cursor = self.conn.execute(""" SELECT block_order FROM document_blocks WHERE block_order > ? ORDER BY block_order LIMIT 1 """, (after_order,)) next_row = cursor.fetchone() if next_row: next_order = next_row['block_order'] gap = next_order - after_order # 3. 检查间隙是否足够 if gap <= 1: # 触发局部重排 self._rebalance_block_orders_between(after_order, next_order) # 重新查询 cursor = self.conn.execute(""" SELECT block_order FROM document_blocks WHERE block_order > ? ORDER BY block_order LIMIT 1 """, (after_order,)) next_row = cursor.fetchone() next_order = next_row['block_order'] if next_row else after_order + 200 # 4. 计算中间值 return (after_order + next_order) // 2 else: # 没有下一个 block,插入到最后 return after_order + 100 else: # 插入到文档末尾 cursor = self.conn.execute(""" SELECT MAX(block_order) as max_order FROM document_blocks """) row = cursor.fetchone() max_order = row['max_order'] if row['max_order'] is not None else 0 return max_order + 100 def _rebalance_indexes_between( self, block_type: str, level: int, start_index: int, end_index: int ): """局部重排:重新分配区间内同类型同级别 blocks 的 index Args: block_type: Block 类型 level: 级别(标题有效) start_index: 起始 index end_index: 结束 index """ # 查询区间内的所有同类型同级别 blocks if block_type == 'heading': cursor = self.conn.execute(""" SELECT id, "index" FROM document_blocks WHERE type = ? AND level = ? AND "index" > ? AND "index" < ? ORDER BY "index" """, (block_type, level, start_index, end_index)) else: cursor = self.conn.execute(""" SELECT id, "index" FROM document_blocks WHERE type = ? AND "index" > ? AND "index" < ? ORDER BY "index" """, (block_type, start_index, end_index)) blocks = cursor.fetchall() if not blocks: return # 计算新的间隔 count = len(blocks) gap = end_index - start_index step = gap // (count + 1) # 重新分配 index new_index = start_index for block in blocks: new_index += step self.conn.execute( 'UPDATE document_blocks SET "index" = ? WHERE id = ?', (new_index, block['id']) ) self.conn.commit() def _rebalance_block_orders_between(self, start_order: int, end_order: int): """局部重排:重新分配区间内所有 blocks 的 block_order Args: start_order: 起始 block_order end_order: 结束 block_order """ # 查询区间内的所有 blocks cursor = self.conn.execute(""" SELECT id, block_order FROM document_blocks WHERE block_order > ? AND block_order < ? ORDER BY block_order """, (start_order, end_order)) blocks = cursor.fetchall() if not blocks: return # 计算新的间隔 count = len(blocks) gap = end_order - start_order step = gap // (count + 1) # 重新分配 block_order new_order = start_order for block in blocks: new_order += step self.conn.execute( 'UPDATE document_blocks SET block_order = ? WHERE id = ?', (new_order, block['id']) ) self.conn.commit() def _row_to_dict(self, row) -> dict: """将 sqlite3.Row 转换为字典""" d = dict(row) # 解析 JSON 字段 if d.get('style'): try: d['style'] = json.loads(d['style']) except (json.JSONDecodeError, TypeError): d['style'] = {} else: d['style'] = {} if d.get('metadata'): try: d['metadata'] = json.loads(d['metadata']) except (json.JSONDecodeError, TypeError): d['metadata'] = {} else: d['metadata'] = {} # content 可能是 JSON(富文本或表格) if d.get('content'): try: d['content'] = json.loads(d['content']) except (json.JSONDecodeError, TypeError): pass # 保持原字符串 return d @staticmethod def generate_block_id(block_type: str, level: int, index: int) -> str: """生成 Block ID Args: block_type: Block 类型 level: 级别(标题有效) index: 序号 Returns: 生成的 Block ID """ if block_type == 'heading': return f'block-h{level}-{index}' elif block_type == 'paragraph': return f'block-p-{index}' elif block_type == 'table': return f'block-table-{index}' elif block_type == 'image': return f'block-img-{index}' elif block_type == 'toc': return f'block-toc-{index}' else: return f'block-{block_type}-{index}'