73 lines
2.2 KiB
Python
73 lines
2.2 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
一键迁移：支持 Docker 或本机执行，需在项目根目录运行。
|
||
python deploy/scripts/migrate.py # 默认用 Docker
|
||
python deploy/scripts/migrate.py --local # 本机执行（需已安装依赖且数据库可连）
|
||
"""
|
||
import os
|
||
import subprocess
|
||
import sys
|
||
import time
|
||
|
||
PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
|
||
BACKEND_DIR = os.path.join(PROJECT_ROOT, "backend")
|
||
|
||
|
||
def run(cmd, cwd=None, env=None):
|
||
r = subprocess.run(cmd, cwd=cwd or PROJECT_ROOT, env=env or os.environ.copy(), shell=True)
|
||
if r.returncode != 0:
|
||
sys.exit(r.returncode)
|
||
|
||
|
||
def migrate_docker():
|
||
os.chdir(PROJECT_ROOT)
|
||
print("[1/3] 启动数据库...")
|
||
run("docker compose up -d db")
|
||
print("[2/3] 等待数据库就绪...")
|
||
for i in range(30):
|
||
r = subprocess.run(
|
||
"docker compose exec -T db pg_isready -U wecom -d wecom_ai",
|
||
shell=True,
|
||
cwd=PROJECT_ROOT,
|
||
capture_output=True,
|
||
)
|
||
if r.returncode == 0:
|
||
break
|
||
time.sleep(1)
|
||
else:
|
||
print("数据库未在 30s 内就绪，请检查 docker compose logs db")
|
||
sys.exit(1)
|
||
print("[3/3] 执行 Alembic 迁移...")
|
||
run('docker compose run --rm backend sh -c "alembic upgrade head"')
|
||
print("迁移完成。")
|
||
|
||
|
||
def migrate_local():
|
||
# 加载 .env
|
||
env_path = os.path.join(PROJECT_ROOT, ".env")
|
||
if os.path.isfile(env_path):
|
||
with open(env_path, encoding="utf-8") as f:
|
||
for line in f:
|
||
line = line.strip()
|
||
if line and not line.startswith("#") and "=" in line:
|
||
k, _, v = line.partition("=")
|
||
os.environ[k.strip()] = v.strip()
|
||
if "DATABASE_URL_SYNC" not in os.environ:
|
||
os.environ.setdefault("DATABASE_URL_SYNC", "postgresql://wecom:wecom_secret@localhost:5432/wecom_ai")
|
||
os.chdir(BACKEND_DIR)
|
||
os.environ["PYTHONPATH"] = BACKEND_DIR
|
||
print("本机执行迁移（backend 目录）...")
|
||
run("alembic upgrade head", cwd=BACKEND_DIR)
|
||
print("迁移完成。")
|
||
|
||
|
||
def main():
|
||
if "--local" in sys.argv:
|
||
migrate_local()
|
||
else:
|
||
migrate_docker()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|