AIyaodianzhushou/data-pipeline/migrate_image_urls.py

#!/usr/bin/env python3
"""
图片 URL 迁移脚本：DB sections/chunks 中的旧图片路径 → 新 hash 文件名路径

同时处理两种可能的 DB 现状（增量安全）：
  - 旧 SQL 已执行:  /images/甲硝唑_rId9.jpg   →  /images/fa123757_rId9.jpg
  - 旧 SQL 未执行:  /images/drugs/甲硝唑_rId9.jpg →  /images/fa123757_rId9.jpg
"""

import re
import sys
import os
import hashlib
import asyncio
from pathlib import Path

import asyncpg


# ======================================================================
# 加载 .env（与环境变量）
# ======================================================================
def _load_env():
    env_file = Path(__file__).resolve().parent.parent / ".env"
    if env_file.exists():
        with open(env_file, encoding="utf-8") as fh:
            for line in fh:
                line = line.strip()
                if not line or line.startswith("#") or "=" not in line:
                    continue
                key, _, val = line.partition("=")
                os.environ.setdefault(key.strip(), val.strip())


_load_env()

# ======================================================================
# 配置
# ======================================================================
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", "pharma2025")

# 新旧 URL 前缀
OLD_PREFIX = "/images/drugs/"
NEW_PREFIX = "/images/"

# 旧文件名安全的规则（与原 docx_ingest.py 的 _safe_drug_name 一致）
_RE_SAFE = re.compile(r'[\s（）()　]+')

# 匹配所有旧格式图片 URL（兼容 /images/ 和 /images/drugs/ 两种前缀）
# group(1) = 旧前缀（中文 safe_name）, group(2) = rId, group(3) = 扩展名
_OLD_IMG_RE = re.compile(
    r'/images/(?:drugs/)?'              # /images/ 或 /images/drugs/
    r'([^"_\s]+)'                        # 旧文件名前缀（中文 safe_name）
    r'_(rId\d+)'                         # _rIdX 引用
    r'\.(jpg|jpeg|png|gif|bmp)'          # 扩展名
)

# 只用于计数查询的 LIKE 模式
_SEARCH_PATTERN = '/images/%_rId%.'


def old_safe_name(drug_name: str) -> str:
    """与原始 docx_ingest.py 的 _safe_drug_name 一致"""
    return _RE_SAFE.sub("_", drug_name).strip("_")


def new_hash_prefix(drug_name: str) -> str:
    """新的 hash 前缀（与更新后的 docx_ingest.py 一致）"""
    return hashlib.md5(drug_name.encode()).hexdigest()[:8]


def replace_urls(text: str, mappings: dict[str, str]) -> str:
    """替换文本中所有旧格式图片 URL → 新 hash 格式。"""

    def _replacer(m: re.Match) -> str:
        old = m.group(1)          # 旧 safe_name（中文）
        rId = m.group(2)         # rIdX
        ext = m.group(3)         # jpg/png/gif
        new = mappings.get(old)
        if new is None:
            # 没匹配到的保留原文件名，但确保前缀是 /images/
            return f"{NEW_PREFIX}{old}_{rId}.{ext}"
        return f"{NEW_PREFIX}{new}_{rId}.{ext}"

    return _OLD_IMG_RE.sub(_replacer, text)


async def main():
    # 1. 连接数据库
    conn = await asyncpg.connect(
        host=DB_HOST, port=DB_PORT, database=DB_NAME,
        user=DB_USER, password=DB_PASSWORD,
    )

    # 2. 查有图片引用的所有药品（同时兼容 /images/drugs/ 和 /images/ 两种状态）
    drugs = await conn.fetch(f"""
        SELECT DISTINCT name FROM drugs
        WHERE sections::text LIKE '%{_SEARCH_PATTERN}%'
    """)
    drug_names = [row["name"] for row in drugs]
    print(f"从 drugs 表找到有图片引用的药品: {len(drug_names)} 个")

    # 3. 构建映射: 旧 safe_name → 新 hash
    mappings: dict[str, str] = {}
    for name in drug_names:
        safe = old_safe_name(name)
        h = new_hash_prefix(name)
        mappings[safe] = h

    print(f"映射条目: {len(mappings)}")
    for i, (old, new) in enumerate(mappings.items()):
        if i >= 10:
            print(f"  ... 还有 {len(mappings) - 10} 条")
            break
        print(f"  {old} → {new}")

    # 4. 统计待更新量
    sections_count = await conn.fetchval(f"""
        SELECT COUNT(*) FROM drugs
        WHERE sections::text LIKE '%{_SEARCH_PATTERN}%'
    """)
    chunks_count = await conn.fetchval(f"""
        SELECT COUNT(*) FROM drug_chunks
        WHERE content LIKE '%{_SEARCH_PATTERN}%'
    """)
    print(f"\n待更新: drugs={sections_count}  chunks={chunks_count}")

    if sections_count == 0 and chunks_count == 0:
        print("✅ 无需更新")
        await conn.close()
        return

    # 5. 更新 drug_chunks.content（纯文本，直接替换）
    if chunks_count > 0:
        print("\n📦 更新 drug_chunks ...")
        rows = await conn.fetch(f"""
            SELECT id, content FROM drug_chunks
            WHERE content LIKE '%{_SEARCH_PATTERN}%'
        """)
        updates = []
        for row in rows:
            new_content = replace_urls(row["content"], mappings)
            if new_content != row["content"]:
                updates.append((new_content, row["id"]))
        if updates:
            async with conn.transaction():
                await conn.executemany(
                    "UPDATE drug_chunks SET content = $1 WHERE id = $2",
                    updates,
                )
            print(f"  ✅ 更新了 {len(updates)} 条")
        else:
            print("  ✅ 无需更新")

    # 6. 更新 drugs.sections（jsonb，需 text ↔ jsonb 转换）
    if sections_count > 0:
        print("\n📦 更新 drugs.sections ...")
        rows = await conn.fetch(f"""
            SELECT drug_id, sections::text AS txt FROM drugs
            WHERE sections::text LIKE '%{_SEARCH_PATTERN}%'
        """)
        updates = []
        for row in rows:
            new_text = replace_urls(row["txt"], mappings)
            if new_text != row["txt"]:
                updates.append((new_text, row["drug_id"]))
        if updates:
            async with conn.transaction():
                await conn.executemany(
                    "UPDATE drugs SET sections = $1::jsonb WHERE drug_id = $2",
                    updates,
                )
            print(f"  ✅ 更新了 {len(updates)} 条")
        else:
            print("  ✅ 无需更新")

    # 7. 验证残留：确认无旧格式 URL
    remaining_d = await conn.fetchval(f"""
        SELECT COUNT(*) FROM drugs
        WHERE sections::text LIKE '%{_SEARCH_PATTERN}%'
    """)
    remaining_c = await conn.fetchval(f"""
        SELECT COUNT(*) FROM drug_chunks
        WHERE content LIKE '%{_SEARCH_PATTERN}%'
    """)

    print(f"\n残留旧格式: drugs={remaining_d}  chunks={remaining_c}")
    if remaining_d == 0 and remaining_c == 0:
        print("✅ 迁移完成")
    else:
        print("⚠️  仍有残留，请手动检查")

    await conn.close()


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