138 lines
3.7 KiB
Python
138 lines
3.7 KiB
Python
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())
|
||
|