mirror of
https://github.com/fawney19/Aether.git
synced 2026-01-02 15:52:26 +08:00
102 lines
2.8 KiB
Python
102 lines
2.8 KiB
Python
"""
|
||
Alembic 环境配置
|
||
用于数据库迁移的运行时环境设置
|
||
"""
|
||
|
||
from logging.config import fileConfig
|
||
from sqlalchemy import engine_from_config, pool
|
||
from alembic import context
|
||
import os
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
# 添加项目根目录到 Python 路径
|
||
sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
|
||
|
||
# 加载 .env 文件（本地开发时需要）
|
||
try:
|
||
from dotenv import load_dotenv
|
||
|
||
env_file = Path(__file__).parent.parent / ".env"
|
||
if env_file.exists():
|
||
load_dotenv(env_file)
|
||
except ImportError:
|
||
pass
|
||
|
||
# 导入所有数据库模型（确保 Alembic 能检测到所有表）
|
||
from src.models.database import Base
|
||
|
||
# Alembic Config 对象
|
||
config = context.config
|
||
|
||
# 从环境变量获取数据库 URL
|
||
# 优先使用 DATABASE_URL，否则从 DB_PASSWORD 自动构建（与 docker-compose 保持一致）
|
||
database_url = os.getenv("DATABASE_URL")
|
||
if not database_url:
|
||
db_password = os.getenv("DB_PASSWORD", "")
|
||
db_host = os.getenv("DB_HOST", "localhost")
|
||
db_port = os.getenv("DB_PORT", "5432")
|
||
db_name = os.getenv("DB_NAME", "aether")
|
||
db_user = os.getenv("DB_USER", "postgres")
|
||
database_url = f"postgresql://{db_user}:{db_password}@{db_host}:{db_port}/{db_name}"
|
||
config.set_main_option("sqlalchemy.url", database_url)
|
||
|
||
# 配置日志
|
||
if config.config_file_name is not None:
|
||
fileConfig(config.config_file_name)
|
||
|
||
# 目标元数据（包含所有表定义）
|
||
target_metadata = Base.metadata
|
||
|
||
|
||
def run_migrations_offline() -> None:
|
||
"""
|
||
离线模式运行迁移
|
||
|
||
在离线模式下，不需要连接数据库，
|
||
只生成 SQL 脚本
|
||
"""
|
||
url = config.get_main_option("sqlalchemy.url")
|
||
context.configure(
|
||
url=url,
|
||
target_metadata=target_metadata,
|
||
literal_binds=True,
|
||
dialect_opts={"paramstyle": "named"},
|
||
compare_type=True, # 比较列类型变更
|
||
compare_server_default=True, # 比较默认值变更
|
||
)
|
||
|
||
with context.begin_transaction():
|
||
context.run_migrations()
|
||
|
||
|
||
def run_migrations_online() -> None:
|
||
"""
|
||
在线模式运行迁移
|
||
|
||
在线模式下，直接连接数据库执行迁移
|
||
"""
|
||
connectable = engine_from_config(
|
||
config.get_section(config.config_ini_section, {}),
|
||
prefix="sqlalchemy.",
|
||
poolclass=pool.NullPool,
|
||
)
|
||
|
||
with connectable.connect() as connection:
|
||
context.configure(
|
||
connection=connection,
|
||
target_metadata=target_metadata,
|
||
compare_type=True, # 比较列类型变更
|
||
compare_server_default=True, # 比较默认值变更
|
||
)
|
||
|
||
with context.begin_transaction():
|
||
context.run_migrations()
|
||
|
||
|
||
# 根据模式选择运行方式
|
||
if context.is_offline_mode():
|
||
run_migrations_offline()
|
||
else:
|
||
run_migrations_online()
|