"""SQLite 内容数据库操作类""" import json import sqlite3 from pathlib import Path from typing import Optional class ContentDB: """SQLite 内容数据库操作类,每个文档对应一个独立的 SQLite 数据库文件""" 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 _build_type_filter(self, block_type: str, level: Optional[int] = None) -> tuple[str, list]: """构建类型过滤条件(用于 index 查询)""" if block_type == 'heading' and level is not None: return "type = ? AND level = ?", [block_type, level] return "type = ?", [block_type] def _query_next_value(self, field: str, after_order: int, type_filter: str, params: list) -> Optional[int]: """查询下一个值(index 或 block_order)""" sql = f'SELECT "{field}" FROM document_blocks WHERE {type_filter} AND block_order > ? ORDER BY block_order LIMIT 1' cursor = self.conn.execute(sql, params + [after_order]) row = cursor.fetchone() return row[field] if row else None def _query_prev_value(self, field: str, after_order: int, type_filter: str, params: list) -> Optional[int]: """查询前一个值(index 或 block_order)""" sql = f'SELECT "{field}" FROM document_blocks WHERE {type_filter} AND block_order <= ? ORDER BY block_order DESC LIMIT 1' cursor = self.conn.execute(sql, params + [after_order]) row = cursor.fetchone() return row[field] if row else None def _query_max_value(self, field: str, type_filter: str, params: list) -> Optional[int]: """查询最大值(index 或 block_order)""" sql = f'SELECT MAX("{field}") as max_val FROM document_blocks' if type_filter: sql += f' WHERE {type_filter}' cursor = self.conn.execute(sql, params) row = cursor.fetchone() return row['max_val'] def _calculate_sparse_value( self, field: str, after_block_id: Optional[str], type_filter: str, params: list, default_min: int, rebalance_func: callable ) -> int: """通用稀疏值计算逻辑""" if after_block_id: 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'] next_val = self._query_next_value(field, after_order, type_filter, params) if next_val is not None: prev_val = self._query_prev_value(field, after_order, type_filter, params) if prev_val is None: prev_val = default_min gap = next_val - prev_val if gap <= 1: rebalance_func(prev_val, next_val) next_val = self._query_next_value(field, after_order, type_filter, params) if next_val is None: next_val = prev_val + 200 return (prev_val + next_val) // 2 else: max_val = self._query_max_value(field, type_filter, params) return (max_val if max_val is not None else default_min) + 100 else: max_val = self._query_max_value(field, type_filter, params) return (max_val if max_val is not None else default_min) + 100 def calculate_index_for_insert(self, after_block_id: str, new_type: str, new_level: int = 0) -> int: """计算插入 block 的 index(稀疏排序)""" type_filter, params = self._build_type_filter(new_type, new_level) return self._calculate_sparse_value( field='index', after_block_id=after_block_id, type_filter=type_filter, params=params, default_min=-100, rebalance_func=lambda prev, next: self._rebalance_indexes_between(new_type, new_level, prev, next) ) def calculate_block_order_for_insert(self, after_block_id: str = None) -> int: """计算插入 block 的 block_order(稀疏排序)""" return self._calculate_sparse_value( field='block_order', after_block_id=after_block_id, type_filter='1=1', params=[], default_min=0, rebalance_func=self._rebalance_block_orders_between ) def _rebalance_indexes_between(self, block_type: str, level: int, start_index: int, end_index: int): """局部重排:重新分配区间内同类型同级别 blocks 的 index""" 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) 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""" 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) 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) 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'] = {} 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""" 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}'