from __future__ import annotations from pathlib import Path import pytest import sqlalchemy as sa from alembic import command from alembic.config import Config from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine _EDGE_INDEXES = { "ix_llmwiki_wiki_edges_knowledge_base_mapping_id", "ix_llmwiki_wiki_edges_source_page_id", "ix_llmwiki_wiki_edges_target_page_id", "ix_llmwiki_wiki_edge_mapping_source", "ix_llmwiki_wiki_edge_mapping_target", } def test_mysql_manual_column_script_keeps_helper_until_all_calls_finish() -> None: sql_path = Path(__file__).parents[1] / "scripts" / "sql" / "column_additions_mysql.sql" sql = sql_path.read_text(encoding="utf-8") drop_offset = sql.index("DROP PROCEDURE IF EXISTS _add_column_if_missing;", sql.index("CREATE PROCEDURE")) call_offsets = [ offset for offset in range(len(sql)) if sql.startswith("CALL _add_column_if_missing(", offset) ] assert call_offsets assert max(call_offsets) < drop_offset assert ( "CALL _add_column_if_missing('llmwiki_wiki_sync_states', " "'last_sync_result', 'TEXT NULL');" ) in sql assert ( "CALL _add_column_if_missing('llmwiki_knowledge_bases', " "'test_questions', 'TEXT NULL');" ) in sql assert ( "CALL _add_column_if_missing('llmwiki_knowledge_bases', " "'retrieval_settings', 'TEXT NULL');" ) in sql assert ( "CALL _add_column_if_missing('llmwiki_knowledge_bases', " "'wiki_parent_directories', 'TEXT NULL');" ) in sql assert ( "CALL _add_column_if_missing('llmwiki_knowledge_bases', " "'wiki_directory_paths', 'TEXT NULL');" ) in sql assert ( "CALL _add_column_if_missing('llmwiki_knowledge_bases', " "'wiki_focus_instructions', 'TEXT NULL');" ) in sql assert ( "CALL _add_column_if_missing('knowledge_retrieval_settings', " "'advanced_management_enabled', 'TINYINT(1) NOT NULL DEFAULT 0');" ) in sql def _config(database: Path) -> Config: migrations = ( Path(__file__).parents[1] / "packages" / "harness" / "deerflow" / "persistence" / "migrations" ) config = Config(str(migrations / "alembic.ini")) config.set_main_option( "sqlalchemy.url", f"sqlite+aiosqlite:///{database.resolve().as_posix()}", ) return config def _prepare_old_database( database: Path, stamp: str, *, has_sync_result: bool, has_retrieval_profile: bool, ) -> None: engine = sa.create_engine(f"sqlite:///{database.resolve().as_posix()}") sync_result = ", last_sync_result JSON" if has_sync_result else "" retrieval_profile = ( ", test_questions JSON, retrieval_settings JSON" if has_retrieval_profile else "" ) with engine.begin() as connection: connection.exec_driver_sql( "CREATE TABLE llmwiki_knowledge_bases " f"(id VARCHAR(36) PRIMARY KEY{retrieval_profile})" ) if has_retrieval_profile: connection.exec_driver_sql( "INSERT INTO llmwiki_knowledge_bases " "(id, test_questions, retrieval_settings) VALUES " "('mapping-1', '[\"保留问题\"]', '{\"top_k\": 8}')" ) connection.exec_driver_sql( "CREATE TABLE llmwiki_wiki_pages (id VARCHAR(36) PRIMARY KEY)" ) connection.exec_driver_sql( "CREATE TABLE llmwiki_wiki_sync_states " f"(id VARCHAR(36) PRIMARY KEY{sync_result})" ) if has_sync_result: connection.exec_driver_sql( "INSERT INTO llmwiki_wiki_sync_states (id, last_sync_result) " "VALUES ('state-1', '{\"processed\": 3}')" ) connection.exec_driver_sql( "CREATE TABLE alembic_version " "(version_num VARCHAR(32) NOT NULL PRIMARY KEY)" ) connection.exec_driver_sql( "INSERT INTO alembic_version (version_num) VALUES (?)", (stamp,), ) engine.dispose() @pytest.mark.parametrize( ("stamp", "has_sync_result", "has_retrieval_profile"), [ ("20260916_01", False, False), ("20260916_02", False, False), ("20260920_01", True, False), ("20260920_01", False, True), ], ) def test_duplicate_historical_stamps_reconcile_to_single_head( tmp_path: Path, stamp: str, has_sync_result: bool, has_retrieval_profile: bool, ) -> None: variant = "sync" if has_sync_result else "profile" if has_retrieval_profile else "base" database = tmp_path / f"old-{stamp}-{variant}.db" _prepare_old_database( database, stamp, has_sync_result=has_sync_result, has_retrieval_profile=has_retrieval_profile, ) config = _config(database) command.upgrade(config, "head") command.upgrade(config, "head") engine = sa.create_engine(f"sqlite:///{database.resolve().as_posix()}") with engine.connect() as connection: inspector = sa.inspect(connection) assert connection.execute( sa.text("SELECT version_num FROM alembic_version") ).scalar_one() == "20261102_02" assert inspector.has_table("task_deeplink_threads") assert connection.execute( sa.text("SELECT COUNT(*) FROM task_deeplink_threads") ).scalar_one() == 0 assert inspector.has_table("task_deeplink_details") assert inspector.has_table("llmwiki_wiki_edges") assert _EDGE_INDEXES <= { str(item["name"]) for item in inspector.get_indexes("llmwiki_wiki_edges") } assert "uq_llmwiki_wiki_edge_mapping_pages_type" in { str(item["name"]) for item in inspector.get_unique_constraints("llmwiki_wiki_edges") } assert len(inspector.get_foreign_keys("llmwiki_wiki_edges")) == 3 assert inspector.has_table("model_capabilities") capability_columns = { str(item["name"]): item for item in inspector.get_columns("model_capabilities") } assert { "model_name", "supports_reasoning_effort", "default_reasoning_effort", "updated_by", "updated_at", } == capability_columns.keys() assert capability_columns["supports_reasoning_effort"]["nullable"] is False mapping_column_info = { str(item["name"]): item for item in inspector.get_columns("llmwiki_knowledge_bases") } assert { "test_questions", "retrieval_settings", "wiki_parent_directories", "wiki_directory_paths", "wiki_focus_instructions", "wiki_content_instructions_raw", "wiki_extraction_instructions_raw", "wiki_generation_managed", "wiki_config_schema_version", "wiki_config_revision", "wiki_config_apply_status", "wiki_config_apply_error", } <= mapping_column_info.keys() assert mapping_column_info["wiki_parent_directories"]["nullable"] is True assert mapping_column_info["wiki_directory_paths"]["nullable"] is True retrieval_setting_columns = { str(item["name"]): item for item in inspector.get_columns("knowledge_retrieval_settings") } assert "advanced_management_enabled" in retrieval_setting_columns assert retrieval_setting_columns["advanced_management_enabled"]["nullable"] is False if has_retrieval_profile: assert connection.execute( sa.text( "SELECT test_questions FROM llmwiki_knowledge_bases " "WHERE id = 'mapping-1'" ) ).scalar_one() == '["保留问题"]' assert connection.execute( sa.text( "SELECT retrieval_settings FROM llmwiki_knowledge_bases " "WHERE id = 'mapping-1'" ) ).scalar_one() == '{"top_k": 8}' sync_columns = { str(item["name"]) for item in inspector.get_columns("llmwiki_wiki_sync_states") } assert "last_sync_result" in sync_columns if has_sync_result: assert connection.execute( sa.text( "SELECT last_sync_result FROM llmwiki_wiki_sync_states " "WHERE id = 'state-1'" ) ).scalar_one() == '{"processed": 3}' engine.dispose() def test_conflicting_20261019_feature_stamp_reconciles_task_schema( tmp_path: Path, ) -> None: """旧知识迁移误用 20261019_01 时，升级仍补齐任务表且不重复加列。""" database = tmp_path / "conflicting-20261019-feature-stamp.db" _prepare_old_database( database, "20260920_01", has_sync_result=True, has_retrieval_profile=True, ) config = _config(database) command.upgrade(config, "20261014_02") engine = sa.create_engine(f"sqlite:///{database.resolve().as_posix()}") with engine.begin() as connection: inspector = sa.inspect(connection) assert inspector.has_table("knowledge_retrieval_settings") assert not inspector.has_table("task_deeplink_threads") connection.exec_driver_sql( "ALTER TABLE knowledge_retrieval_settings " "ADD COLUMN advanced_management_enabled BOOLEAN NOT NULL DEFAULT 0" ) connection.exec_driver_sql( "INSERT INTO knowledge_retrieval_settings " "(scope, smart_reasoning_enabled, advanced_management_enabled, " "revision, updated_at) VALUES " "('global', 1, 1, 2, '2026-10-19 00:00:00')" ) connection.exec_driver_sql("DELETE FROM alembic_version") connection.exec_driver_sql( "INSERT INTO alembic_version (version_num) VALUES ('20261019_01')" ) engine.dispose() command.upgrade(config, "head") command.upgrade(config, "head") engine = sa.create_engine(f"sqlite:///{database.resolve().as_posix()}") with engine.connect() as connection: inspector = sa.inspect(connection) assert connection.execute( sa.text("SELECT version_num FROM alembic_version") ).scalar_one() == "20261102_02" retrieval_columns = [ str(item["name"]) for item in inspector.get_columns("knowledge_retrieval_settings") ] assert retrieval_columns.count("advanced_management_enabled") == 1 assert connection.execute( sa.text( "SELECT advanced_management_enabled " "FROM knowledge_retrieval_settings WHERE scope = 'global'" ) ).scalar_one() in (1, True) assert inspector.has_table("task_deeplink_threads") task_columns = { str(item["name"]) for item in inspector.get_columns("task_deeplink_threads") } assert {"sub_type", "display_name", "status", "metadata_json"} <= task_columns task_indexes = { str(item["name"]): [str(column) for column in item["column_names"]] for item in inspector.get_indexes("task_deeplink_threads") } assert task_indexes["ix_task_deeplink_threads_scope"] == [ "task_id", "go_path", "sub_type", "updated_at", "thread_id", ] assert task_indexes["ix_task_deeplink_threads_task_go_updated"] == [ "task_id", "go_path", "updated_at", "thread_id", ] engine.dispose() def test_broken_two_column_detail_key_is_repaired_at_head(tmp_path: Path) -> None: """20261022_02 的中间版本只加了列没换成三列主键（batch API 不存在）。 这类库 alembic_version 已到 head 但键仍是两列，typed upsert 必炸 UNIQUE 约束。兜底修复迁移应把键换成三列且保留旧数据。用最小库直接 从坏状态升级（不跑全链，链上其它迁移依赖 create_all 的表）。 """ database = tmp_path / "broken-detail-key.db" engine = sa.create_engine(f"sqlite:///{database.resolve().as_posix()}") try: with engine.begin() as connection: connection.exec_driver_sql( "CREATE TABLE alembic_version (version_num VARCHAR(32) NOT NULL PRIMARY KEY)" ) connection.exec_driver_sql( "INSERT INTO alembic_version (version_num) VALUES ('20261023_02')" ) # 「半执行」现场：sub_type 列存在，主键仍是两列，有一行旧快照。 connection.exec_driver_sql( "CREATE TABLE task_deeplink_details (" "detail_type VARCHAR(32) NOT NULL, " "detail_id VARCHAR(128) NOT NULL, " "task_name VARCHAR(512) NOT NULL, " "task_content TEXT NOT NULL, " "overview TEXT NOT NULL, " "extra_json JSON NOT NULL, " "version INTEGER NOT NULL, " "created_by VARCHAR(64), " "updated_by VARCHAR(64), " "created_at DATETIME NOT NULL, " "updated_at DATETIME NOT NULL, " "sub_type VARCHAR(64) NOT NULL DEFAULT '', " "PRIMARY KEY (detail_type, detail_id))" ) connection.exec_driver_sql( "INSERT INTO task_deeplink_details " "(detail_type, detail_id, task_name, task_content, overview, " "extra_json, version, created_at, updated_at, sub_type) VALUES " "('b5-1', '001', '旧任务', '旧详情', '', '{}', 1, " "CURRENT_TIMESTAMP, CURRENT_TIMESTAMP, '')" ) finally: engine.dispose() command.upgrade(_config(database), "head") engine = sa.create_engine(f"sqlite:///{database.resolve().as_posix()}") try: with engine.connect() as connection: inspector = sa.inspect(connection) pk = [ str(column) for column in inspector.get_pk_constraint("task_deeplink_details")[ "constrained_columns" ] ] assert pk == ["detail_type", "detail_id", "sub_type"] assert connection.execute( sa.text( "SELECT detail_type, detail_id, sub_type, task_name " "FROM task_deeplink_details" ) ).fetchall() == [("b5-1", "001", "", "旧任务")] finally: engine.dispose() def test_custom_business_mapping_migration_preserves_existing_rows(tmp_path: Path) -> None: """The new type columns and widened history key upgrade a populated install.""" database = tmp_path / "custom-business-mapping.db" engine = sa.create_engine(f"sqlite:///{database.resolve().as_posix()}") try: with engine.begin() as connection: connection.exec_driver_sql( "CREATE TABLE alembic_version (version_num VARCHAR(32) NOT NULL PRIMARY KEY)" ) connection.exec_driver_sql( "INSERT INTO alembic_version (version_num) VALUES ('20261102_01')" ) connection.exec_driver_sql( "CREATE TABLE business_mapping (" "business_code VARCHAR(32) PRIMARY KEY, label VARCHAR(128) NOT NULL, " "agent_id VARCHAR(64), agent_name VARCHAR(191), chain_id VARCHAR(64), " "chain_title VARCHAR(512), sort_order INTEGER NOT NULL, " "created_at DATETIME NOT NULL, updated_at DATETIME NOT NULL)" ) connection.exec_driver_sql( "INSERT INTO business_mapping " "(business_code, label, sort_order, created_at, updated_at) VALUES " "('3Q', '3QFX', 0, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)" ) connection.exec_driver_sql( "CREATE TABLE task_deeplink_threads (" "thread_id VARCHAR(64) PRIMARY KEY, task_id VARCHAR(128) NOT NULL, " "source_task_id VARCHAR(128) NOT NULL, go_path VARCHAR(16) NOT NULL, " "sub_type VARCHAR(64) NOT NULL DEFAULT '', agent_id VARCHAR(128), " "display_name VARCHAR(256), status VARCHAR(20) NOT NULL DEFAULT 'idle', " "metadata_json JSON, created_by VARCHAR(64), " "created_at DATETIME NOT NULL, updated_at DATETIME NOT NULL)" ) connection.exec_driver_sql( "INSERT INTO task_deeplink_threads " "(thread_id, task_id, source_task_id, go_path, created_at, updated_at) VALUES " "('thread-1', 'task-1', 'task-1', '3qfx', CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)" ) command.upgrade(_config(database), "head") with engine.connect() as connection: inspector = sa.inspect(connection) mapping_columns = { str(column["name"]): column for column in inspector.get_columns("business_mapping") } assert {"mapping_type", "is_custom"} <= mapping_columns.keys() assert connection.execute( sa.text( "SELECT mapping_type, is_custom FROM business_mapping " "WHERE business_code = '3Q'" ) ).one() == ("内置业务", 0) assert connection.execute( sa.text("SELECT go_path FROM task_deeplink_threads WHERE thread_id = 'thread-1'") ).scalar_one() == "3qfx" assert connection.execute( sa.text("SELECT version_num FROM alembic_version") ).scalar_one() == "20261102_02" finally: engine.dispose() async def _exercise_parent_directories(store) -> str: configured = ["专题资料", "政策法规"] directory_paths = [["专题资料"], ["专题资料", "国内政策"]] retrieval_settings = { "entity_expansion_enabled": True, "max_related_terms": 6, "statistics_enabled": True, } created = await store.create_mapping( weknora_id="remote-parent-directories", owner_user_id="user-one", name="知识库", description="", kb_type="document", retrieval_settings=retrieval_settings, wiki_parent_directories=configured, wiki_directory_paths=directory_paths, wiki_focus_instructions="重点提取适用对象", wiki_content_instructions_raw="保留关键数字", wiki_extraction_instructions_raw="抽取时间线", wiki_generation_managed=True, wiki_config_schema_version=2, wiki_config_revision=1, wiki_config_apply_status="active", ) mapping_id = str(created["id"]) assert created["wiki_parent_directories"] == configured assert all(isinstance(item, str) for item in created["wiki_parent_directories"]) assert created["wiki_directory_paths"] == directory_paths assert created["wiki_focus_instructions"] == "重点提取适用对象" assert created["wiki_generation_managed"] is True assert created["wiki_config_revision"] == 1 assert created["retrieval_settings"] == retrieval_settings configured.append("调用方后续修改") directory_paths[1].append("调用方后续修改") retrieval_settings["max_related_terms"] = 12 unchanged = await store.update_mapping(mapping_id, description="已更新") assert unchanged is not None assert unchanged["wiki_parent_directories"] == ["专题资料", "政策法规"] assert unchanged["wiki_directory_paths"] == [ ["专题资料"], ["专题资料", "国内政策"], ] assert unchanged["retrieval_settings"]["max_related_terms"] == 6 emptied = await store.update_mapping(mapping_id, wiki_parent_directories=[]) assert emptied is not None assert emptied["wiki_parent_directories"] == [] reset = await store.update_mapping( mapping_id, wiki_parent_directories=["重新设置"], wiki_directory_paths=[["重新设置"], ["重新设置", "二级"]], wiki_focus_instructions="新重点", wiki_config_revision=2, ) assert reset is not None assert reset["wiki_parent_directories"] == ["重新设置"] assert reset["wiki_directory_paths"] == [["重新设置"], ["重新设置", "二级"]] assert reset["wiki_focus_instructions"] == "新重点" assert reset["wiki_config_revision"] == 2 cleared = await store.update_mapping(mapping_id, wiki_parent_directories=None) assert cleared is not None assert cleared["wiki_parent_directories"] is None return mapping_id @pytest.mark.asyncio async def test_memory_llmwiki_parent_directories_persistence() -> None: from deerflow.persistence.llmwiki.memory import MemoryLlmWikiStore store = MemoryLlmWikiStore() mapping_id = await _exercise_parent_directories(store) saved = await store.get_authorized( mapping_id, "user-one", write=True, is_admin=False, ) assert saved is not None assert saved["wiki_parent_directories"] is None @pytest.mark.asyncio async def test_sql_llmwiki_parent_directories_persistence(tmp_path: Path) -> None: from deerflow.persistence.base import Base from deerflow.persistence.llmwiki.model import LlmWikiKnowledgeBaseRow from deerflow.persistence.llmwiki.sql import LlmWikiRepository engine = create_async_engine( f"sqlite+aiosqlite:///{tmp_path / 'llmwiki-parent-directories.db'}" ) async with engine.begin() as connection: await connection.run_sync( lambda sync_connection: Base.metadata.create_all( sync_connection, tables=[LlmWikiKnowledgeBaseRow.__table__], ) ) session_factory = async_sessionmaker(engine, expire_on_commit=False) store = LlmWikiRepository(session_factory) mapping_id = await _exercise_parent_directories(store) async with session_factory() as session: persisted = await session.scalar( sa.select(LlmWikiKnowledgeBaseRow.wiki_parent_directories).where( LlmWikiKnowledgeBaseRow.id == mapping_id ) ) assert persisted is None await engine.dispose()