life-echo/api/scripts/run_memoir_images_migration.py

"""
将 chapters.cover_image 与 chapter_sections.image 的 JSON 数据迁移到 memoir_images 表（字段独立列）。

前置：先执行 api/migrations/add_memoir_images_table.sql 建表。
用法（在项目根目录或 api 目录下）:
    python -m api.scripts.run_memoir_images_migration
"""
import json
import logging
import os
import sys
import uuid

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
os.chdir(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from dotenv import load_dotenv
load_dotenv()

from sqlalchemy import text
from sqlalchemy.engine import Engine

logging.basicConfig(level=logging.INFO, format="%(message)s")
logger = logging.getLogger(__name__)


def get_engine() -> Engine:
    from sqlalchemy import create_engine
    url = os.getenv("DATABASE_URL", "postgresql://postgres:postgres@localhost:5432/life_echo")
    return create_engine(url.replace("postgresql://", "postgresql+psycopg://"), pool_pre_ping=True)


def _row_from_image_json(img: dict | None, chapter_id: str, section_id: str | None, order_index: int) -> dict | None:
    if not img or not isinstance(img, dict):
        return None
    placeholder = (img.get("placeholder") or "").strip()
    description = (img.get("description") or "").strip()
    if not placeholder and not description:
        return None
    if not placeholder:
        placeholder = f'{{{{IMAGE:{description}}}}}'
    created = img.get("created_at")
    updated = img.get("updated_at")
    if isinstance(created, str) and created:
        try:
            from datetime import datetime
            created = datetime.fromisoformat(created.replace("Z", "+00:00"))
        except Exception:
            created = None
    if isinstance(updated, str) and updated:
        try:
            from datetime import datetime
            updated = datetime.fromisoformat(updated.replace("Z", "+00:00"))
        except Exception:
            updated = None
    return {
        "id": str(uuid.uuid4()).replace("-", "")[:32],
        "chapter_id": chapter_id,
        "section_id": section_id,
        "order_index": order_index,
        "placeholder": placeholder or None,
        "description": description or None,
        "status": (img.get("status") or "pending").strip() or "pending",
        "prompt": img.get("prompt") or None,
        "url": img.get("url") or None,
        "storage_key": img.get("storage_key") or None,
        "provider": img.get("provider") or None,
        "style": img.get("style") or None,
        "size": img.get("size") or None,
        "error": img.get("error") or None,
        "retryable": img.get("retryable") if img.get("retryable") is not None else None,
        "created_at": created,
        "updated_at": updated,
    }


def run_sql_migration(engine: Engine):
    from pathlib import Path
    sql_path = Path(__file__).parent.parent / "migrations" / "add_memoir_images_table.sql"
    if not sql_path.exists():
        logger.warning("未找到 %s，请先执行该 SQL 建表", sql_path)
        return
    sql = sql_path.read_text(encoding="utf-8")
    stmts = []
    rest = sql
    while rest:
        rest = rest.lstrip()
        if rest.startswith("--"):
            rest = rest[rest.find("\n") + 1:] if "\n" in rest else ""
            continue
        if rest.upper().startswith("DO "):
            i = rest.find("$$")
            if i == -1:
                break
            j = rest.find("$$", i + 2)
            if j == -1:
                break
            stmts.append(rest[: j + 2].strip() + ";")
            rest = rest[j + 2:].lstrip().lstrip(";").lstrip()
            continue
        idx = rest.find(";")
        if idx == -1:
            break
        part = rest[: idx].strip()
        rest = rest[idx + 1:]
        if part and not part.startswith("--"):
            stmts.append(part + ";")
    with engine.begin() as conn:
        for i, s in enumerate(stmts):
            try:
                conn.execute(text(s))
                logger.info("  SQL %s OK", i + 1)
            except Exception as e:
                if "already exists" in str(e).lower():
                    logger.info("  SQL %s (已存在)", i + 1)
                    continue
                raise
    logger.info("1/2 SQL 迁移完成")


def run_data_migration(engine: Engine):
    ins = text("""
        INSERT INTO memoir_images (
            id, chapter_id, section_id, order_index,
            placeholder, description, status, prompt, url, storage_key,
            provider, style, size, error, retryable, created_at, updated_at
        ) VALUES (
            :id, :chapter_id, :section_id, :order_index,
            :placeholder, :description, :status, :prompt, :url, :storage_key,
            :provider, :style, :size, :error, :retryable, :created_at, :updated_at
        )
    """)
    with engine.connect() as conn:
        r = conn.execute(text("""
            SELECT id, cover_image FROM chapters WHERE cover_image IS NOT NULL
        """))
        cover_count = 0
        for row in r:
            ch_id, cover = row[0], row[1]
            if isinstance(cover, str):
                try:
                    cover = json.loads(cover)
                except Exception:
                    cover = None
            if not cover or not isinstance(cover, dict):
                continue
            exists = conn.execute(
                text("SELECT 1 FROM memoir_images WHERE chapter_id = :ch_id AND section_id IS NULL"),
                {"ch_id": ch_id},
            ).fetchone()
            if exists:
                continue
            row_data = _row_from_image_json(cover, ch_id, None, 0)
            if not row_data:
                continue
            conn.execute(ins, {**row_data, "updated_at": row_data.get("updated_at")})
            conn.commit()
            cover_count += 1
        logger.info("封面图迁移: %d 条", cover_count)

    with engine.connect() as conn:
        r = conn.execute(text("""
            SELECT id, chapter_id, order_index, image FROM chapter_sections WHERE image IS NOT NULL
        """))
        sec_count = 0
        for row in r:
            sec_id, ch_id, ord_idx, img = row[0], row[1], row[2], row[3]
            if isinstance(img, str):
                try:
                    img = json.loads(img)
                except Exception:
                    img = None
            if not img or not isinstance(img, dict):
                continue
            exists = conn.execute(
                text("SELECT 1 FROM memoir_images WHERE section_id = :sec_id"),
                {"sec_id": sec_id},
            ).fetchone()
            if exists:
                continue
            row_data = _row_from_image_json(img, ch_id, sec_id, ord_idx + 1)
            if not row_data:
                continue
            conn.execute(ins, {**row_data, "updated_at": row_data.get("updated_at")})
            conn.commit()
            sec_count += 1
        logger.info("段落配图迁移: %d 条", sec_count)
    logger.info("2/2 数据迁移完成")


if __name__ == "__main__":
    logger.info("开始 memoir_images 迁移…")
    engine = get_engine()
    run_sql_migration(engine)
    run_data_migration(engine)
    logger.info("迁移全部完成")