From 0a3e430d08c589a750d6e6a81ebf15bf76c5a9e2 Mon Sep 17 00:00:00 2001 From: edy Date: Tue, 28 Jul 2026 11:24:52 +0800 Subject: [PATCH] =?UTF-8?q?ES=E5=90=8C=E6=AD=A5:=20=E6=94=B9=E4=B8=BA?= =?UTF-8?q?=E5=85=88=E6=B8=85=E7=A9=BA=E7=B4=A2=E5=BC=95=E5=86=8D=E6=8C=89?= =?UTF-8?q?=20Mongo=20=E5=85=A8=E9=87=8F=E9=87=8D=E5=BB=BA=E5=86=99?= =?UTF-8?q?=E5=85=A5?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Cursor --- scripts/sync_es_from_comm.py | 186 ++++++++++++++------------------ tests/test_sync_es_from_comm.py | 132 +++++++++++++++-------- 2 files changed, 172 insertions(+), 146 deletions(-) diff --git a/scripts/sync_es_from_comm.py b/scripts/sync_es_from_comm.py index a160799..0fcb000 100644 --- a/scripts/sync_es_from_comm.py +++ b/scripts/sync_es_from_comm.py @@ -1,10 +1,10 @@ #!/usr/bin/env python3 -"""Mongo Somni 集合 → ES 差异同步（单文件：适配、备份、对账、向量化、定时调度）。 +"""Mongo Somni 集合 → ES 全量重建同步（单文件：适配、备份、清空、向量化、定时调度）。 -以 Mongo _id 为准，只读不写源库： - - ES 有、Mongo 无 → 删 ES - - Mongo 有 → 比差异，有变才 upsert - - 先同步 somni_audio_tag_dictionary（name/name_en 向量），再同步 somni_audio_materials +以 Mongo `_id` 为准，只读不写源库： + - 先删除目标 ES 索引全部数据（删索引再 ensure 重建） + - 再按 Mongo 启用文档全量插入；ES 文档 id = Mongo `_id`，_source 不含 `id`/`_id` + - 先同步 somni_audio_tag_dictionary，再同步 somni_audio_materials 服务启动后按 SYNC_INTERVAL_DAYS 注册定时任务；也可手动执行本脚本。 @@ -91,53 +91,49 @@ def mongo_doc_id(doc: dict[str, Any]) -> str: def material_doc_to_es(doc: dict[str, Any]) -> dict[str, Any] | None: - """Mongo 原料文档 → ES 文档（去掉 _id）；无 _id 或无 audio_url 则跳过。""" + """Mongo 原料文档 → ES _source（去掉 _id/id）；无 _id 或无 audio_url 则跳过。""" doc_id = mongo_doc_id(doc) if not doc_id: return None - return material_source_for_es(bson_to_jsonable(doc)) - - -def tag_dictionary_compare_snapshot(doc: dict[str, Any]) -> dict[str, Any]: - """标签词典 diff 快照（不含向量）。""" - keys = ( - "type", - "code", - "status", - "name", - "name_en", - "description", - "applicability", - "parent_tag_id", - "created_at", - "updated_at", - "created_by", - "updated_by", - ) - return {k: doc.get(k) for k in keys} - - -def material_compare_snapshot(doc: dict[str, Any]) -> dict[str, Any]: - """原料 diff 快照（不比较 dense vector，避免浮点噪声导致反复更新）。""" - snapshot = bson_to_jsonable(doc) - snapshot.pop("_id", None) - snapshot.pop("id", None) - snapshot.pop("description_vector", None) - return snapshot - - -def tag_documents_differ(desired: dict[str, Any], existing: dict[str, Any]) -> bool: - return tag_dictionary_compare_snapshot(desired) != tag_dictionary_compare_snapshot(existing) + payload = material_source_for_es(bson_to_jsonable(doc)) + if payload is None: + return None + payload.pop("_id", None) + payload.pop("id", None) + return payload -def material_documents_differ(desired: dict[str, Any], existing: dict[str, Any]) -> bool: - return material_compare_snapshot(desired) != material_compare_snapshot(existing) +def tag_doc_to_es(doc: dict[str, Any]) -> dict[str, Any] | None: + """Mongo 标签文档 → ES _source（去掉 _id/id）；无 _id 则跳过。""" + doc_id = mongo_doc_id(doc) + if not doc_id: + return None + payload = bson_to_jsonable(doc) + payload.pop("_id", None) + payload.pop("id", None) + return payload def zero_vector(dim: int) -> list[float]: return [0.0] * dim +async def wipe_and_recreate_index( + es_client: AsyncElasticsearch, + es_search: EsSearch, + index: str, +) -> int: + """删除索引内全部数据：统计文档数 → 删索引 → ensure 重建映射。""" + count = 0 + if await es_client.indices.exists(index=index): + count_resp = await es_client.count(index=index) + count = int(count_resp.get("count", 0)) + await es_client.indices.delete(index=index) + logger.info("已删除 ES 索引以全量重建：{}，原文档数={}", index, count) + await es_search.ensure_indices() + return count + + def write_backup(path: Path, records: list[dict[str, Any]]) -> None: if path.is_file(): path.unlink() @@ -226,39 +222,33 @@ async def run(self, *, dry_run: bool) -> dict[str, int]: if not dry_run: write_backup(self._settings.sync_tag_dictionary_backup_path, bson_to_jsonable(docs)) - source_ids = {mongo_doc_id(d) for d in docs if mongo_doc_id(d)} - es_ids = await self._es_search.list_all_tag_dictionary_doc_ids() - - for doc_id in es_ids - source_ids: - if dry_run: - stats["deleted"] += 1 - continue - try: - await self._client.delete(index=self._es_search.tag_dictionary_index, id=doc_id) - stats["deleted"] += 1 - except Exception as exc: + payloads: dict[str, dict[str, Any]] = {} + for doc in docs: + es_doc = tag_doc_to_es(doc) + if es_doc is None: stats["failed"] += 1 - logger.error("删除 ES 孤儿标签失败，id={}，原因：{}", doc_id, exc) + continue + payloads[mongo_doc_id(doc)] = es_doc - for doc in docs: - outcome = await self._sync_one(doc, dry_run=dry_run) + if dry_run: + es_ids = await self._es_search.list_all_tag_dictionary_doc_ids() + stats["deleted"] = len(es_ids) + stats["created"] = len(payloads) + return stats + + stats["deleted"] = await wipe_and_recreate_index( + self._client, + self._es_search, + self._es_search.tag_dictionary_index, + ) + for doc_id, es_doc in payloads.items(): + outcome = await self._insert_one(doc_id, es_doc) stats[outcome] += 1 - if not dry_run: - self._es_search.clear_content_tag_vectors_cache() + self._es_search.clear_content_tag_vectors_cache() return stats - async def _sync_one(self, doc: dict[str, Any], *, dry_run: bool) -> str: - doc_id = mongo_doc_id(doc) - if not doc_id: - return "failed" - es_doc = bson_to_jsonable(doc) - es_doc.pop("_id", None) - existing = await self._es_search.get_tag_dictionary_source(doc_id) - if existing and not tag_documents_differ(es_doc, existing): - return "unchanged" - if dry_run: - return "created" if existing is None else "updated" + async def _insert_one(self, doc_id: str, es_doc: dict[str, Any]) -> str: try: es_doc["name_vector"] = await self._encoder.encode_one(str(es_doc.get("name", ""))) name_en = str(es_doc.get("name_en", "")).strip() @@ -270,7 +260,7 @@ async def _sync_one(self, doc: dict[str, Any], *, dry_run: bool) -> str: id=doc_id, document=es_doc, ) - return "created" if existing is None else "updated" + return "created" except Exception as exc: logger.error( "同步标签词典失败，id={}，name={}，原因：{}", @@ -321,40 +311,34 @@ async def run(self, *, dry_run: bool) -> dict[str, int]: continue payloads[mongo_doc_id(doc)] = es_doc - es_ids = await self._es_search.list_all_audio_doc_ids() - for doc_id in es_ids - set(payloads.keys()): - if dry_run: - stats["deleted"] += 1 - continue - try: - await self._client.delete(index=self._es_search.audio_index, id=doc_id) - stats["deleted"] += 1 - except Exception as exc: - stats["failed"] += 1 - logger.error("删除 ES 孤儿原料失败，id={}，原因：{}", doc_id, exc) + if dry_run: + es_ids = await self._es_search.list_all_audio_doc_ids() + stats["deleted"] = len(es_ids) + stats["created"] = len(payloads) + return stats + stats["deleted"] = await wipe_and_recreate_index( + self._client, + self._es_search, + self._es_search.audio_index, + ) for doc_id, es_doc in payloads.items(): - outcome = await self._sync_one(doc_id, es_doc, dry_run=dry_run) + outcome = await self._insert_one(doc_id, es_doc) stats[outcome] += 1 return stats - async def _sync_one(self, doc_id: str, es_doc: dict[str, Any], *, dry_run: bool) -> str: - existing = await self._es_search.get_audio_source(doc_id) - if ( - existing - and not material_documents_differ(es_doc, existing) - and existing.get("description_vector") - ): - return "unchanged" - if dry_run: - return "created" if existing is None else "updated" + async def _insert_one(self, doc_id: str, es_doc: dict[str, Any]) -> str: try: es_doc["description_vector"] = await self._encoder.encode_one( str(es_doc.get("description_text", "")) ) - await self._client.index(index=self._es_search.audio_index, id=doc_id, document=es_doc) - return "created" if existing is None else "updated" + await self._client.index( + index=self._es_search.audio_index, + id=doc_id, + document=es_doc, + ) + return "created" except Exception as exc: logger.error( "同步原料失败，id={}，name={}，原因：{}", @@ -383,7 +367,7 @@ def __init__( self._settings = settings async def run(self, *, dry_run: bool = False) -> SyncJobResult: - logger.info("开始 Mongo → ES 差异同步，dry_run={}", dry_run) + logger.info("开始 Mongo → ES 全量重建同步，dry_run={}", dry_run) await self._es_search.migrate_legacy_indices() await self._es_search.ensure_indices() @@ -417,20 +401,16 @@ async def run(self, *, dry_run: bool = False) -> SyncJobResult: material_failed=material_stats["failed"], ) logger.info( - "Mongo → ES 同步结束：标签 拉取={} 删={} 增={} 改={} 未变={} 失败={}；" - "原料 拉取={} 跳过={} 删={} 增={} 改={} 未变={} 失败={} dry_run={}", + "Mongo → ES 全量同步结束：标签 拉取={} 清索引={} 增={} 失败={}；" + "原料 拉取={} 跳过={} 清索引={} 增={} 失败={} dry_run={}", result.tag_fetched, result.tag_deleted, result.tag_created, - result.tag_updated, - result.tag_unchanged, result.tag_failed, result.material_fetched, result.material_skipped, result.material_deleted, result.material_created, - result.material_updated, - result.material_unchanged, result.material_failed, dry_run, ) @@ -482,7 +462,7 @@ def start_sync_scheduler(state: AppState, settings: Settings) -> None: return async def _job() -> None: - logger.info("定时任务触发：Mongo → ES 差异同步") + logger.info("定时任务触发：Mongo → ES 全量重建同步") await run_scheduled_sync(state, settings) _scheduler = AsyncIOScheduler(timezone=UTC) @@ -547,8 +527,8 @@ async def _load_stage(stage: str) -> list: def main() -> None: - parser = argparse.ArgumentParser(description="Mongo Somni 集合差异同步至 ES") - parser.add_argument("--dry-run", action="store_true", help="只拉取比对，不写 ES、不备份") + parser = argparse.ArgumentParser(description="Mongo Somni 集合全量重建同步至 ES") + parser.add_argument("--dry-run", action="store_true", help="只拉取统计，不删 ES、不写 ES、不备份") args = parser.parse_args() exit_code = asyncio.run(_run_cli(dry_run=args.dry_run)) if exit_code != 0: diff --git a/tests/test_sync_es_from_comm.py b/tests/test_sync_es_from_comm.py index e48fa57..f621721 100644 --- a/tests/test_sync_es_from_comm.py +++ b/tests/test_sync_es_from_comm.py @@ -1,4 +1,4 @@ -"""scripts/sync_es_from_comm.py Mongo 同步逻辑单元测试。""" +"""scripts/sync_es_from_comm.py Mongo 全量重建同步单元测试。""" from __future__ import annotations @@ -17,11 +17,10 @@ _redact_mongo_uri, bson_to_jsonable, material_doc_to_es, - material_documents_differ, mongo_doc_id, start_sync_scheduler, - tag_dictionary_compare_snapshot, - tag_documents_differ, + tag_doc_to_es, + wipe_and_recreate_index, zero_vector, ) @@ -42,6 +41,7 @@ def _material_doc( ) -> dict: return { "_id": ObjectId(doc_id) if len(doc_id) == 24 else doc_id, + "id": doc_id, "audio_name": audio_name, "description": "描述", "status": True, @@ -63,6 +63,7 @@ def _material_doc( def _tag_doc(doc_id: str, *, name: str = "放松", name_en: str = "Unwind") -> dict: return { "_id": ObjectId(doc_id) if len(doc_id) == 24 else doc_id, + "id": doc_id, "type": "sleep_stage", "code": "unwind", "status": "启用", @@ -87,38 +88,52 @@ def test_material_doc_to_es_requires_audio_url() -> None: assert material_doc_to_es(doc) is None -def test_tag_documents_same_when_only_vectors_differ() -> None: - desired = {"name": "放松", "name_en": "Unwind", "status": "启用"} - existing = { - "name": "放松", - "name_en": "Unwind", - "status": "启用", - "name_vector": [0.1], - "name_en_vector": [0.2], - } - assert tag_documents_differ(desired, existing) is False +def test_material_doc_to_es_keeps_sleep_stage_names_without_id_fields() -> None: + payload = material_doc_to_es(_material_doc("6a33a7928030d4cf420efeb6")) + assert payload is not None + assert "id" not in payload + assert "_id" not in payload + assert payload["sleep_stage_names"] == ["放松"] + + +def test_tag_doc_to_es_strips_id_fields() -> None: + payload = tag_doc_to_es(_tag_doc("6a325acc1a3dbc128504c423")) + assert payload is not None + assert "id" not in payload + assert "_id" not in payload + assert payload["name"] == "放松" + + +@pytest.mark.asyncio +async def test_wipe_and_recreate_index_deletes_then_ensures() -> None: + es_client = MagicMock() + es_client.indices.exists = AsyncMock(return_value=True) + es_client.count = AsyncMock(return_value={"count": 12}) + es_client.indices.delete = AsyncMock() + es_search = MagicMock() + es_search.ensure_indices = AsyncMock() + deleted = await wipe_and_recreate_index(es_client, es_search, "somni_audio_materials") -def test_material_documents_differ_on_tag_change() -> None: - desired = material_doc_to_es(_material_doc("6a33a7928030d4cf420efeb6")) - existing = dict(desired) - existing["sleep_stage_tags"] = [] - assert material_documents_differ(desired, existing) is True + assert deleted == 12 + es_client.indices.delete.assert_awaited_once_with(index="somni_audio_materials") + es_search.ensure_indices.assert_awaited_once() @pytest.mark.asyncio -async def test_tag_sync_job_deletes_es_orphan() -> None: +async def test_tag_sync_job_wipes_index_then_inserts() -> None: mongo = MagicMock() mongo.fetch_tag_dictionary = AsyncMock(return_value=[_tag_doc("6a325acc1a3dbc128504c423")]) es_search = MagicMock() - es_search.list_all_tag_dictionary_doc_ids = AsyncMock( - return_value={"6a325acc1a3dbc128504c423", "orphan"} - ) - es_search.get_tag_dictionary_source = AsyncMock(return_value=None) + es_search.list_all_tag_dictionary_doc_ids = AsyncMock(return_value={"orphan"}) es_search.tag_dictionary_index = "somni_audio_tag_dictionary" + es_search.clear_content_tag_vectors_cache = MagicMock() + es_search.ensure_indices = AsyncMock() es_client = MagicMock() + es_client.indices.exists = AsyncMock(return_value=True) + es_client.count = AsyncMock(return_value={"count": 1}) + es_client.indices.delete = AsyncMock() es_client.index = AsyncMock() - es_client.delete = AsyncMock() encoder = MagicMock() encoder.encode_one = AsyncMock(return_value=[0.1] * 512) @@ -127,23 +142,27 @@ async def test_tag_sync_job_deletes_es_orphan() -> None: ).run(dry_run=False) assert stats["deleted"] == 1 - es_client.delete.assert_awaited_once_with( - index="somni_audio_tag_dictionary", id="orphan" - ) + assert stats["created"] == 1 + es_client.indices.delete.assert_awaited_once_with(index="somni_audio_tag_dictionary") + kwargs = es_client.index.await_args.kwargs + assert kwargs["id"] == "6a325acc1a3dbc128504c423" + assert "id" not in kwargs["document"] + assert "_id" not in kwargs["document"] @pytest.mark.asyncio -async def test_material_sync_job_skips_unchanged(tmp_path) -> None: +async def test_material_sync_job_wipes_then_reindexes(tmp_path) -> None: doc = _material_doc("6a33a7928030d4cf420efeb6") - es_payload = material_doc_to_es(doc) - es_payload["description_vector"] = [0.1] * 512 mongo = MagicMock() mongo.fetch_materials = AsyncMock(return_value=[doc]) es_search = MagicMock() es_search.list_all_audio_doc_ids = AsyncMock(return_value={"6a33a7928030d4cf420efeb6"}) - es_search.get_audio_source = AsyncMock(return_value=es_payload) es_search.audio_index = "somni_audio_materials" + es_search.ensure_indices = AsyncMock() es_client = MagicMock() + es_client.indices.exists = AsyncMock(return_value=True) + es_client.count = AsyncMock(return_value={"count": 99}) + es_client.indices.delete = AsyncMock() es_client.index = AsyncMock() encoder = MagicMock() encoder.encode_one = AsyncMock(return_value=[0.1] * 512) @@ -152,8 +171,14 @@ async def test_material_sync_job_skips_unchanged(tmp_path) -> None: mongo, es_search, es_client, encoder, Settings(sync_backup_dir=str(tmp_path)) ).run(dry_run=False) - assert stats["unchanged"] == 1 - es_client.index.assert_not_called() + assert stats["deleted"] == 99 + assert stats["created"] == 1 + es_client.indices.delete.assert_awaited_once_with(index="somni_audio_materials") + indexed = es_client.index.await_args.kwargs + assert indexed["id"] == "6a33a7928030d4cf420efeb6" + assert "id" not in indexed["document"] + assert "_id" not in indexed["document"] + assert indexed["document"]["sleep_stage_names"] == ["放松"] @pytest.mark.asyncio @@ -163,9 +188,11 @@ async def test_material_sync_job_writes_description_vector(tmp_path) -> None: mongo.fetch_materials = AsyncMock(return_value=[doc]) es_search = MagicMock() es_search.list_all_audio_doc_ids = AsyncMock(return_value=set()) - es_search.get_audio_source = AsyncMock(return_value=None) es_search.audio_index = "somni_audio_materials" + es_search.ensure_indices = AsyncMock() es_client = MagicMock() + es_client.indices.exists = AsyncMock(return_value=False) + es_client.indices.delete = AsyncMock() es_client.index = AsyncMock() encoder = MagicMock() encoder.encode_one = AsyncMock(return_value=[0.2] * 512) @@ -181,6 +208,29 @@ async def test_material_sync_job_writes_description_vector(tmp_path) -> None: encoder.encode_one.assert_awaited_once_with(indexed["description_text"]) +@pytest.mark.asyncio +async def test_material_sync_dry_run_does_not_wipe(tmp_path) -> None: + doc = _material_doc("6a33a7928030d4cf420efeb6") + mongo = MagicMock() + mongo.fetch_materials = AsyncMock(return_value=[doc]) + es_search = MagicMock() + es_search.list_all_audio_doc_ids = AsyncMock(return_value={"a", "b", "c"}) + es_search.audio_index = "somni_audio_materials" + es_client = MagicMock() + es_client.indices.delete = AsyncMock() + es_client.index = AsyncMock() + encoder = MagicMock() + + stats = await MaterialsSyncJob( + mongo, es_search, es_client, encoder, Settings(sync_backup_dir=str(tmp_path)) + ).run(dry_run=True) + + assert stats["deleted"] == 3 + assert stats["created"] == 1 + es_client.indices.delete.assert_not_called() + es_client.index.assert_not_called() + + @pytest.mark.asyncio async def test_mongo_sync_job_migrates_legacy_indices() -> None: mongo = MagicMock() @@ -191,7 +241,11 @@ async def test_mongo_sync_job_migrates_legacy_indices() -> None: es_search.ensure_indices = AsyncMock() es_search.list_all_tag_dictionary_doc_ids = AsyncMock(return_value=set()) es_search.list_all_audio_doc_ids = AsyncMock(return_value=set()) + es_search.tag_dictionary_index = "somni_audio_tag_dictionary" + es_search.audio_index = "somni_audio_materials" + es_search.clear_content_tag_vectors_cache = MagicMock() es_client = MagicMock() + es_client.indices.exists = AsyncMock(return_value=False) encoder = MagicMock() await MongoEsSyncJob(mongo, es_search, es_client, encoder, Settings()).run(dry_run=True) @@ -205,14 +259,6 @@ def test_zero_vector_has_embedding_dim_length() -> None: assert all(v == 0.0 for v in zero_vector(512)) -def test_tag_dictionary_snapshot_uses_created_by_fields() -> None: - doc = _tag_doc("6a325acc1a3dbc128504c423") - snapshot = tag_dictionary_compare_snapshot(bson_to_jsonable(doc)) - assert snapshot["created_by"] == "tester" - assert snapshot["updated_by"] == "tester" - assert "create_by" not in snapshot - - def test_start_sync_scheduler_skipped_without_mongo_uri() -> None: with patch("scripts.sync_es_from_comm.AsyncIOScheduler") as mock_cls: settings = Settings(sync_enabled=True, mongo_uri="")