"""MySQL async client wrapper for persistent storage. Provides connection-pool-based helpers for research sessions, research records, convert logs, and general queries. """ from __future__ import annotations from typing import Any try: import aiomysql except ModuleNotFoundError: # pragma: no cover - exercised only when dependency is absent aiomysql = None from app.config import settings class MySQLClient: """Thin async wrapper around an aiomysql connection pool.""" def __init__(self, pool: Any) -> None: self._pool = pool @classmethod async def from_settings(cls) -> "MySQLClient": if aiomysql is None: raise RuntimeError("aiomysql is not installed") pool = await aiomysql.create_pool( host=settings.mysql_host, port=settings.mysql_port, user=settings.mysql_user, password=settings.mysql_password, db=settings.mysql_database, charset=settings.mysql_charset, autocommit=True, connect_timeout=settings.mysql_connect_timeout, minsize=settings.mysql_pool_minsize, maxsize=settings.mysql_pool_maxsize, ) return cls(pool) async def close(self) -> None: self._pool.close() await self._pool.wait_closed() async def ensure_research_sessions_table(self) -> None: await self.execute( """ CREATE TABLE IF NOT EXISTS research_sessions ( session_id VARCHAR(128) NOT NULL, user_id VARCHAR(128) NOT NULL, history JSON NOT NULL, created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, expires_at DATETIME NULL, PRIMARY KEY (session_id, user_id), KEY idx_research_sessions_expires_at (expires_at), KEY idx_research_sessions_updated_at (updated_at) ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci """ ) async def ensure_convert_log_table(self) -> None: await self.execute( """ CREATE TABLE IF NOT EXISTS convert_log ( id BIGINT UNSIGNED NOT NULL AUTO_INCREMENT, doc_id VARCHAR(256) NOT NULL COMMENT '文档 ID', task_id VARCHAR(128) DEFAULT NULL COMMENT 'Celery 任务 ID', source_format VARCHAR(32) NOT NULL COMMENT '源文件格式 (doc/xls/ppt/wps/et/ofd)', target_format VARCHAR(32) NOT NULL COMMENT '目标格式 (docx/xlsx/pptx/pdf)', source_path VARCHAR(1024) DEFAULT NULL COMMENT '源文件路径', output_path VARCHAR(1024) DEFAULT NULL COMMENT '转换后文件路径', success TINYINT(1) NOT NULL DEFAULT 0 COMMENT '是否成功 (0=失败 1=成功)', duration_ms INT DEFAULT NULL COMMENT '转换耗时(毫秒)', file_size BIGINT DEFAULT NULL COMMENT '源文件大小(字节)', error_message TEXT DEFAULT NULL COMMENT '失败时的错误信息', created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP COMMENT '创建时间', PRIMARY KEY (id), KEY idx_convert_log_doc_id (doc_id), KEY idx_convert_log_task_id (task_id), KEY idx_convert_log_created_at (created_at) ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci COMMENT='文档格式转换日志' """ ) # ------------------------------------------------------------------ # Research records tables # ------------------------------------------------------------------ async def ensure_research_records_table(self) -> None: await self.execute( """ CREATE TABLE IF NOT EXISTS research_records ( id VARCHAR(128) NOT NULL, user_id VARCHAR(128) NOT NULL, session_id VARCHAR(128) DEFAULT NULL COMMENT '关联运行时会话（Phase 2 扩展位）', title VARCHAR(500) NOT NULL, mode VARCHAR(32) NOT NULL DEFAULT 'deep', status VARCHAR(32) NOT NULL DEFAULT 'draft', output_template VARCHAR(64) NOT NULL DEFAULT 'comprehensive', summary TEXT DEFAULT NULL COMMENT '列表摘要', archived TINYINT(1) NOT NULL DEFAULT 0, task_json JSON DEFAULT NULL, plan_json JSON DEFAULT NULL, references_json JSON DEFAULT NULL, imported_items_json JSON DEFAULT NULL, clarification_messages_json JSON DEFAULT NULL COMMENT '澄清对话历史', chat_messages_json JSON DEFAULT NULL COMMENT '研究后对话历史', notes TEXT DEFAULT NULL, parent_record_id VARCHAR(128) DEFAULT NULL COMMENT '父版本记录 ID', root_record_id VARCHAR(128) DEFAULT NULL COMMENT '根记录 ID（版本链）', version_no INT NOT NULL DEFAULT 1, last_error TEXT DEFAULT NULL, created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, completed_at DATETIME DEFAULT NULL, PRIMARY KEY (id), KEY idx_research_records_user_id (user_id), KEY idx_research_records_status (status), KEY idx_research_records_root_record_id (root_record_id), KEY idx_research_records_updated_at (updated_at), UNIQUE KEY uq_root_version (root_record_id, version_no) ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci COMMENT='研究记录主表' """ ) async def ensure_research_record_reports_table(self) -> None: await self.execute( """ CREATE TABLE IF NOT EXISTS research_record_reports ( record_id VARCHAR(128) NOT NULL, report_json JSON NOT NULL, final_document_md LONGTEXT DEFAULT NULL COMMENT '交付物文档（完整 Markdown）', created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, PRIMARY KEY (record_id) ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci COMMENT='研究报告子表（大字段独立存储）' """ ) # Migration: add final_document_md column to existing tables try: await self.execute( "ALTER TABLE research_record_reports ADD COLUMN final_document_md LONGTEXT DEFAULT NULL COMMENT '交付物文档（完整 Markdown）'" ) except Exception as exc: # Ignore "Duplicate column name" (MySQL error 1060), re-raise others if "1060" not in str(exc) and "Duplicate column" not in str(exc): raise async def ensure_research_record_runs_table(self) -> None: await self.execute( """ CREATE TABLE IF NOT EXISTS research_record_runs ( run_id VARCHAR(128) NOT NULL, record_id VARCHAR(128) NOT NULL, user_id VARCHAR(128) NOT NULL, run_type VARCHAR(32) NOT NULL COMMENT 'full 或 section_rerun', status VARCHAR(32) NOT NULL DEFAULT 'pending', section_title VARCHAR(500) DEFAULT NULL COMMENT '仅 section_rerun 时有值', started_at DATETIME DEFAULT NULL, completed_at DATETIME DEFAULT NULL, error TEXT DEFAULT NULL, created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, PRIMARY KEY (run_id), KEY idx_research_record_runs_record_id (record_id), KEY idx_research_record_runs_user_id (user_id), KEY idx_research_record_runs_status (status), KEY idx_research_record_runs_created_at (created_at) ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci COMMENT='研究执行实例表' """ ) # ------------------------------------------------------------------ # Notebook tables # ------------------------------------------------------------------ async def ensure_notebook_tables(self) -> None: await self.execute( """ CREATE TABLE IF NOT EXISTS notebooks ( id VARCHAR(128) NOT NULL PRIMARY KEY, user_id VARCHAR(128) NOT NULL, title VARCHAR(500) NOT NULL DEFAULT '未命名 Notebook', description TEXT, status VARCHAR(32) NOT NULL DEFAULT 'active', config_json JSON, created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, INDEX idx_nb_user (user_id), INDEX idx_nb_status (status) ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci COMMENT='Notebook 主表' """ ) await self.execute( """ CREATE TABLE IF NOT EXISTS notebook_sources ( id VARCHAR(128) NOT NULL PRIMARY KEY, notebook_id VARCHAR(128) NOT NULL, user_id VARCHAR(128) NOT NULL, source_type VARCHAR(32) NOT NULL, title VARCHAR(500) NOT NULL DEFAULT '', doc_id VARCHAR(256), content_hash VARCHAR(128), file_path VARCHAR(1024), paste_text LONGTEXT, ingest_status VARCHAR(32) NOT NULL DEFAULT 'pending', ingest_task_id VARCHAR(128), ingest_error TEXT, metadata_json JSON, summary TEXT, selected TINYINT(1) NOT NULL DEFAULT 1, sort_order INT NOT NULL DEFAULT 0, created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, INDEX idx_ns_notebook (notebook_id), UNIQUE KEY uk_ns_nb_doc (notebook_id, doc_id) ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci COMMENT='Notebook 来源表' """ ) await self.execute( """ CREATE TABLE IF NOT EXISTS notebook_messages ( id VARCHAR(128) NOT NULL PRIMARY KEY, notebook_id VARCHAR(128) NOT NULL, user_id VARCHAR(128) NOT NULL, session_id VARCHAR(128) NOT NULL DEFAULT 'default', role VARCHAR(16) NOT NULL, content LONGTEXT, references_json JSON, graph_data_json JSON, suggestions_json JSON, created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, INDEX idx_nm_nb_sess (notebook_id, session_id), INDEX idx_nm_created (created_at) ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci COMMENT='Notebook 对话消息表' """ ) await self.execute( """ CREATE TABLE IF NOT EXISTS notebook_outputs ( id VARCHAR(128) NOT NULL PRIMARY KEY, notebook_id VARCHAR(128) NOT NULL, user_id VARCHAR(128) NOT NULL, output_type VARCHAR(64) NOT NULL, title VARCHAR(500) NOT NULL DEFAULT '', content_md LONGTEXT, status VARCHAR(32) NOT NULL DEFAULT 'generating', error TEXT, task_id VARCHAR(128), context_json JSON, created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, INDEX idx_no_notebook (notebook_id) ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci COMMENT='Notebook 输出文档表' """ ) # ------------------------------------------------------------------ # Query helpers # ------------------------------------------------------------------ async def fetch_one( self, query: str, params: tuple[Any, ...] | None = None, ) -> dict[str, Any] | None: if aiomysql is None: raise RuntimeError("aiomysql is not installed") async with self._pool.acquire() as conn: async with conn.cursor(aiomysql.DictCursor) as cursor: await cursor.execute(query, params or ()) row = await cursor.fetchone() return dict(row) if row else None async def fetch_all( self, query: str, params: tuple[Any, ...] | None = None, ) -> list[dict[str, Any]]: if aiomysql is None: raise RuntimeError("aiomysql is not installed") async with self._pool.acquire() as conn: async with conn.cursor(aiomysql.DictCursor) as cursor: await cursor.execute(query, params or ()) rows = await cursor.fetchall() return [dict(row) for row in rows] if rows else [] async def execute( self, query: str, params: tuple[Any, ...] | None = None, ) -> int: async with self._pool.acquire() as conn: async with conn.cursor() as cursor: affected = await cursor.execute(query, params or ()) return int(affected or 0)