ax-backend-v1/app/services/content_db.py

"""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 _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