shaguabijia-app-server/scripts/init_postgres.py

"""Bootstrap PostgreSQL: 建用户 + 建库 + 授权 + 写 .env + 跑迁移。

新机器初始化用。前置:已装 PostgreSQL 16 + 知道 postgres 超级用户密码。

用法:
    python scripts/init_postgres.py

环境变量(可选,不填会交互式问):
    PG_HOST          PG 主机, 默认 localhost
    PG_PORT          PG 端口, 默认 5432
    PG_SUPER_USER    超级用户, 默认 postgres
    PG_SUPER_PASS    超级用户密码 (不填会 getpass 交互输入)
    APP_DB_NAME      业务库名, 默认 shaguabijia
    APP_DB_USER      业务用户名, 默认 shaguabijia_app
    APP_DB_PASS      业务用户密码 (不填会自动生成强密码)

幂等:已存在的用户/库不会重复建,已配的 .env 会被覆盖前提示。
"""
from __future__ import annotations

import getpass
import os
import secrets
import subprocess
import sys
from pathlib import Path

try:
    import psycopg
    from psycopg import sql
except ImportError:
    print("❌ 缺 psycopg。先跑: pip install -e .")
    sys.exit(1)


ROOT = Path(__file__).resolve().parent.parent
ENV_FILE = ROOT / ".env"
ENV_EXAMPLE = ROOT / ".env.example"


def env_or_input(key: str, prompt: str, default: str = "", secret: bool = False) -> str:
    val = os.environ.get(key, "").strip()
    if val:
        return val
    if secret:
        return getpass.getpass(f"{prompt}: ").strip()
    suffix = f" [{default}]" if default else ""
    raw = input(f"{prompt}{suffix}: ").strip()
    return raw or default


def ensure_role(conn: psycopg.Connection, user: str, password: str) -> None:
    # 注意: CREATE/ALTER USER 不支持参数化密码, 必须用 sql.Literal 内联
    with conn.cursor() as cur:
        cur.execute("SELECT 1 FROM pg_roles WHERE rolname=%s", (user,))
        if cur.fetchone():
            print(f"  · 用户 {user} 已存在, 重置密码")
            cur.execute(
                sql.SQL("ALTER USER {} WITH PASSWORD {}").format(
                    sql.Identifier(user), sql.Literal(password)
                )
            )
        else:
            print(f"  · 创建用户 {user}")
            cur.execute(
                sql.SQL("CREATE USER {} WITH PASSWORD {}").format(
                    sql.Identifier(user), sql.Literal(password)
                )
            )


def ensure_database(conn: psycopg.Connection, db: str, owner: str) -> None:
    with conn.cursor() as cur:
        cur.execute("SELECT 1 FROM pg_database WHERE datname=%s", (db,))
        if cur.fetchone():
            print(f"  · 数据库 {db} 已存在")
            return
        print(f"  · 创建数据库 {db} (owner={owner}, encoding=UTF8)")
        cur.execute(
            sql.SQL("CREATE DATABASE {} OWNER {} ENCODING 'UTF8'").format(
                sql.Identifier(db), sql.Identifier(owner)
            )
        )


def grant_all(conn: psycopg.Connection, db: str, user: str) -> None:
    with conn.cursor() as cur:
        cur.execute(
            sql.SQL("GRANT ALL PRIVILEGES ON DATABASE {} TO {}").format(
                sql.Identifier(db), sql.Identifier(user)
            )
        )


def write_env(database_url: str) -> None:
    if not ENV_FILE.exists():
        if ENV_EXAMPLE.exists():
            ENV_FILE.write_text(ENV_EXAMPLE.read_text(encoding="utf-8"), encoding="utf-8")
            print(f"  · 从 .env.example 复制出 .env")
        else:
            ENV_FILE.write_text("", encoding="utf-8")

    lines = ENV_FILE.read_text(encoding="utf-8").splitlines()
    new_lines = []
    replaced = False
    for line in lines:
        if line.startswith("DATABASE_URL="):
            old_url = line.split("=", 1)[1]
            if old_url.strip() and old_url.strip() != database_url:
                ans = input(f"\n.env 里已有 DATABASE_URL=\n  {old_url}\n覆盖吗? [y/N]: ").strip().lower()
                if ans != "y":
                    print("  · 保留原 DATABASE_URL")
                    new_lines.append(line)
                    replaced = True
                    continue
            new_lines.append(f"DATABASE_URL={database_url}")
            replaced = True
        else:
            new_lines.append(line)
    if not replaced:
        new_lines.append(f"DATABASE_URL={database_url}")

    ENV_FILE.write_text("\n".join(new_lines) + "\n", encoding="utf-8")
    print(f"  · 写入 .env -> DATABASE_URL")


def run_alembic_upgrade() -> bool:
    print("\n[3/3] 跑 alembic upgrade head")
    try:
        subprocess.run(
            [sys.executable, "-m", "alembic", "upgrade", "head"],
            cwd=ROOT,
            check=True,
        )
        return True
    except subprocess.CalledProcessError as e:
        print(f"❌ alembic 失败: {e}")
        return False


def main() -> int:
    print("=" * 60)
    print("PostgreSQL 初始化脚本 — shaguabijia-app-server")
    print("=" * 60)

    host = env_or_input("PG_HOST", "PG host", "localhost")
    port = env_or_input("PG_PORT", "PG port", "5432")
    super_user = env_or_input("PG_SUPER_USER", "PG superuser", "postgres")
    super_pass = env_or_input("PG_SUPER_PASS", "PG superuser password", secret=True)
    db_name = env_or_input("APP_DB_NAME", "App database name", "shaguabijia")
    db_user = env_or_input("APP_DB_USER", "App database user", "shaguabijia_app")
    db_pass = os.environ.get("APP_DB_PASS", "").strip()
    if not db_pass:
        db_pass = secrets.token_urlsafe(32)
        print(f"  · 自动生成业务用户密码: {db_pass}")

    print(f"\n[1/3] 连接 PG (host={host}:{port}, user={super_user})")
    try:
        conn = psycopg.connect(
            host=host, port=int(port), user=super_user, password=super_pass,
            dbname="postgres", autocommit=True,
        )
    except psycopg.OperationalError as e:
        print(f"❌ 连不上 PG: {e}")
        print("   检查 1) PG 服务是否在跑 2) 超级用户密码是否正确 3) 端口防火墙是否放行")
        return 1

    print("\n[2/3] 建用户 + 建库 + 授权")
    with conn:
        ensure_role(conn, db_user, db_pass)
        ensure_database(conn, db_name, db_user)
        grant_all(conn, db_name, db_user)

    database_url = f"postgresql+psycopg://{db_user}:{db_pass}@{host}:{port}/{db_name}"
    write_env(database_url)

    if not run_alembic_upgrade():
        return 1

    print("\n" + "=" * 60)
    print("✅ 全部完成")
    print(f"   DATABASE_URL = {database_url}")
    print("   下一步: 启动服务 -> ./run.sh 或 uvicorn app.main:app --reload --port 8770")
    print("=" * 60)
    return 0


if __name__ == "__main__":
    sys.exit(main())