from __future__ import annotations import asyncio import os import sys from logging.config import fileConfig from pathlib import Path from alembic import context from sqlalchemy import pool from sqlalchemy.engine import Connection from sqlalchemy.ext.asyncio import async_engine_from_config # 让 alembic 在 `server/` 下运行时也能 import app.* SERVER_DIR = Path(__file__).resolve().parents[1] # .../server/alembic -> .../server sys.path.append(str(SERVER_DIR)) from app.db.base import Base # noqa: E402 import app.db.models # noqa: F401,E402 # 确保模型被导入,metadata 完整 # Alembic Config 对象 config = context.config # 配置日志 if config.config_file_name is not None: fileConfig(config.config_file_name) # 目标 metadata(autogenerate 依赖) target_metadata = Base.metadata def _read_env_kv(env_path: Path) -> dict[str, str]: """ 读取 .env 文件中的 KEY=VALUE。 说明: - 迁移阶段只需要 DATABASE_URL,不应因为 Redis/Celery 等配置缺失而失败 - 这里不依赖 pydantic-settings 的 Settings 校验,避免“缺字段导致迁移不可用” """ data: dict[str, str] = {} if not env_path.exists(): return data for raw in env_path.read_text(encoding="utf-8").splitlines(): line = raw.strip() if not line or line.startswith("#"): continue if "=" not in line: continue k, v = line.split("=", 1) k = k.strip() v = v.strip().strip('"').strip("'") if k: data[k] = v return data def _get_database_url() -> str: """ 获取数据库连接串。 约定: - 优先读取环境变量 `DATABASE_URL` - 若未设置,则按 `APP_ENV`(默认 dev)读取 `server/.env.dev` 或 `server/.env.prod` 注意:迁移阶段仅依赖 DATABASE_URL;不应强制要求 REDIS_URL / CELERY_BROKER_URL 等配置存在。 """ # 允许在 alembic 命令时临时覆盖 env_url = os.getenv("DATABASE_URL") if env_url: return env_url app_env = (os.getenv("APP_ENV") or "dev").strip() or "dev" env_file = SERVER_DIR / (".env.prod" if app_env == "prod" else ".env.dev") kv = _read_env_kv(env_file) url = kv.get("DATABASE_URL") if url: return url raise RuntimeError( "缺少 DATABASE_URL:请设置环境变量 DATABASE_URL,或在 server/.env.dev(或 .env.prod)中配置 DATABASE_URL。" ) def run_migrations_offline() -> None: """离线模式:生成 SQL 脚本,不连接数据库。""" url = _get_database_url() context.configure( url=url, target_metadata=target_metadata, literal_binds=True, dialect_opts={"paramstyle": "named"}, compare_type=True, ) with context.begin_transaction(): context.run_migrations() def do_run_migrations(connection: Connection) -> None: """在线模式:在已有连接上执行迁移。""" context.configure( connection=connection, target_metadata=target_metadata, compare_type=True, ) with context.begin_transaction(): context.run_migrations() async def run_migrations_online() -> None: """在线模式:使用异步引擎执行迁移。""" url = _get_database_url() config.set_main_option("sqlalchemy.url", url) connectable = async_engine_from_config( config.get_section(config.config_ini_section) or {}, prefix="sqlalchemy.", poolclass=pool.NullPool, ) async with connectable.connect() as connection: await connection.run_sync(do_run_migrations) await connectable.dispose() if context.is_offline_mode(): run_migrations_offline() else: asyncio.run(run_migrations_online())