Maas-Model-platform/backend/scripts/run_043_parsed_pricing_migration.py

#!/usr/bin/env python3
"""
仅执行 043_create_parsed_pricing_tables.sql

用法：
  python run_043_parsed_pricing_migration.py
  python run_043_parsed_pricing_migration.py --schema aigcspace
"""
import sys
from pathlib import Path

from sqlalchemy import create_engine, text
from sqlalchemy.exc import SQLAlchemyError
from dotenv import load_dotenv

# 确保可导入项目模块
PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from app.database import DATABASE_URL


def load_env() -> None:
    env_path = PROJECT_ROOT / ".env"
    load_dotenv(dotenv_path=env_path)


def main() -> None:
    import argparse

    load_env()

    parser = argparse.ArgumentParser(description="执行 043_create_parsed_pricing_tables.sql")
    parser.add_argument("--schema", default="aigcspace", help="schema 名称")
    args = parser.parse_args()

    migration_path = PROJECT_ROOT / "migrations" / "043_create_parsed_pricing_tables.sql"
    if not migration_path.exists():
        raise SystemExit(f"未找到迁移文件: {migration_path}")

    sql = migration_path.read_text(encoding="utf-8")
    if not sql.strip():
        print("迁移文件为空，跳过。")
        return

    engine = create_engine(DATABASE_URL)
    try:
        with engine.begin() as conn:
            conn.execute(text(f"CREATE SCHEMA IF NOT EXISTS {args.schema}"))
            conn.execute(text(sql))
        print("迁移执行完成。")
    except SQLAlchemyError as exc:
        print(f"迁移执行失败: {exc}")
        raise SystemExit(1)


if __name__ == "__main__":
    main()