AIyaodianzhushou/data-pipeline/verify_import.py

"""
导入完整性验证脚本
检查 DOCX 文件数 vs 数据库入库数，定位漏读/漏写的药品
"""
import os, sys, json, asyncio
from pathlib import Path
from collections import Counter

# 添加同目录模块
sys.path.insert(0, str(Path(__file__).resolve().parent))
from docx_ingest import find_docx_files, parse_docx

# 数据库
import asyncpg

SOURCE_DIR = os.environ.get("DOCX_SOURCE_DIR", "/opt/2025")
DB_HOST = os.environ.get("POSTGRES_HOST", "localhost")
DB_PORT = os.environ.get("POSTGRES_PORT", "5432")
DB_NAME = os.environ.get("POSTGRES_DB", "pharmacopoeia")
DB_USER = os.environ.get("POSTGRES_USER", "postgres")
DB_PASSWORD = os.environ.get("POSTGRES_PASSWORD", "postgres")


async def main():
    # ==========================================
    # 1. 扫描 DOCX 文件
    # ==========================================
    print("=" * 60)
    print("📂 扫描 DOCX 文件...")
    files = find_docx_files(SOURCE_DIR)
    print(f"   DOCX 文件总数: {len(files)}")

    # 提取文件名（去 .docx）作为药品名
    docx_names = set()
    docx_by_volume = Counter()
    for fp in files:
        name = Path(fp).stem
        docx_names.add(name)
        vol = "一部"
        for k in ["output4", "output3", "output2", "output"]:
            if k in fp:
                vol = {"output": "一部", "output2": "二部", "output3": "三部", "output4": "四部"}[k]
                break
        docx_by_volume[vol] += 1
    print(f"   去重后唯一药名: {len(docx_names)}")
    print(f"   各卷分布: {dict(docx_by_volume)}")

    # ==========================================
    # 2. 查询数据库
    # ==========================================
    print("\n📊 查询数据库...")
    conn = await asyncpg.connect(
        host=DB_HOST, port=DB_PORT, database=DB_NAME,
        user=DB_USER, password=DB_PASSWORD
    )

    # 2.1 药品统计
    total_drugs = await conn.fetchval("SELECT COUNT(*) FROM drugs WHERE is_active = TRUE")
    total_chunks = await conn.fetchval("SELECT COUNT(*) FROM drug_chunks WHERE vec IS NOT NULL")
    total_chunks_null = await conn.fetchval("SELECT COUNT(*) FROM drug_chunks WHERE vec IS NULL")

    print(f"   drugs 表 (active): {total_drugs}")
    print(f"   drug_chunks (有向量): {total_chunks}")
    print(f"   drug_chunks (无向量): {total_chunks_null}")

    # 2.2 获取数据库中所有药名
    db_names = set()
    rows = await conn.fetch("SELECT name FROM drugs WHERE is_active = TRUE")
    for r in rows:
        db_names.add(r["name"])
    print(f"   数据库中药名: {len(db_names)}")

    # 2.3 按来源版本统计
    version_counts = await conn.fetch("""
        SELECT source_version, COUNT(*) FROM drugs
        WHERE is_active = TRUE GROUP BY source_version ORDER BY COUNT(*) DESC
    """)
    print(f"   版本分布:")
    for r in version_counts:
        print(f"     {r['source_version']}: {r['count']}")

    # 2.4 按分类统计
    cat_counts = await conn.fetch("""
        SELECT category, COUNT(*) FROM drugs
        WHERE is_active = TRUE GROUP BY category ORDER BY COUNT(*) DESC
    """)
    print(f"   分类分布:")
    for r in cat_counts:
        print(f"     {r['category']}: {r['count']}")

    # ==========================================
    # 3. 对比分析
    # ==========================================
    print("\n🔍 对比分析:")

    # 3.1 DOCX 有但 DB 没有的（漏读）
    missing = docx_names - db_names
    if missing:
        print(f"\n   ❌ DOCX 有但数据库没有（{len(missing)} 个）:")
        for name in sorted(missing)[:30]:
            print(f"      - {name}")
        if len(missing) > 30:
            print(f"      ... 还有 {len(missing) - 30} 个")
    else:
        print(f"\n   ✅ 无漏读：所有 DOCX 药名均已在数据库中")

    # 3.2 DB 有但 DOCX 没有的（额外数据，如 Wiki）
    extra = db_names - docx_names
    if extra:
        print(f"\n   📝 数据库额外药名（Wiki/样本数据，{len(extra)} 个）:")
        for name in sorted(extra):
            print(f"      - {name}")

    # ==========================================
    # 4. 抽查特定药品的 sections
    # ==========================================
    print("\n🔬 抽查药品内容完整性:")
    spot_check = ["布洛芬", "甲硝唑", "阿莫西林", "对乙酰氨基酚", "二甲双胍"]
    for name in spot_check:
        drug = await conn.fetchrow(
            "SELECT drug_id, name, sections, source_version FROM drugs WHERE name = $1 AND is_active = TRUE",
            name
        )
        if drug:
            sections = json.loads(drug["sections"]) if drug["sections"] else {}
            chunk_count = await conn.fetchval(
                "SELECT COUNT(*) FROM drug_chunks WHERE drug_id = $1 AND vec IS NOT NULL",
                drug["drug_id"]
            )
            sec_names = list(sections.keys())
            print(f"   ✅ {name} | {len(sec_names)} sections | {chunk_count} chunks | {drug['source_version']}")
            print(f"      栏目: {sec_names[:10]}...")
        else:
            print(f"   ❌ {name}: 数据库中不存在")

    # ==========================================
    # 5. 检查 chunks 完整性
    # ==========================================
    print("\n📐 Chunks 完整性检查:")

    # 5.1 检查 drug_id 引用完整性（孤儿 chunk）
    orphan_chunks = await conn.fetchval("""
        SELECT COUNT(*) FROM drug_chunks c
        LEFT JOIN drugs d ON d.drug_id = c.drug_id
        WHERE d.id IS NULL
    """)
    print(f"   孤儿 chunks（无对应 drug）: {orphan_chunks}")

    # 5.2 有 drug 但没有 chunk 的
    no_chunk_drugs = await conn.fetchval("""
        SELECT COUNT(*) FROM drugs d
        WHERE d.is_active = TRUE
        AND NOT EXISTS (SELECT 1 FROM drug_chunks c WHERE c.drug_id = d.drug_id)
    """)
    print(f"   无 chunk 的药品: {no_chunk_drugs}")

    # 5.3 vec 维度是否正确
    vec_dim = await conn.fetchval("""
        SELECT vector_dims(vec) FROM drug_chunks WHERE vec IS NOT NULL LIMIT 1
    """)
    print(f"   向量维度: {vec_dim}")

    # ==========================================
    # 6. 汇总结论
    # ==========================================
    print("\n" + "=" * 60)
    print("📋 验证结论:")
    print(f"   DOCX 文件: {len(files)}")
    print(f"   数据库中药品: {total_drugs}")
    print(f"   数据库 chunks: {total_chunks}")
    print(f"   漏读药品: {len(missing)}")
    print(f"   额外药品 (Wiki等): {len(extra)}")
    print(f"   孤儿 chunks: {orphan_chunks}")
    print(f"   无 chunk 药品: {no_chunk_drugs}")
    if len(missing) == 0 and orphan_chunks == 0 and total_chunks > 0:
        print("\n   ✅ 导入完整，无漏读漏写")
    else:
        print(f"\n   ⚠️  发现问题，请检查上述 ❌ 项")
    print("=" * 60)

    await conn.close()


if __name__ == "__main__":
    asyncio.run(main())