57 lines
2.0 KiB
Python
57 lines
2.0 KiB
Python
"""初始化 mcp_auth 库 + 建表（一次性脚本）。
|
||
|
||
用法：
|
||
# 开发库（47.101.220.40）
|
||
python seed.py --host 47.101.220.40 --user postgres --password digiwin
|
||
|
||
# 生产库（10.100.154.100）
|
||
python seed.py --host 10.100.154.100 --user postgres --password digiwin
|
||
|
||
幂等：可重复执行，表已存在则跳过。
|
||
"""
|
||
|
||
import argparse
|
||
import asyncio
|
||
import os
|
||
from pathlib import Path
|
||
|
||
import asyncpg
|
||
|
||
INIT_SQL_PATH = Path(__file__).resolve().parent.parent / "sql" / "init.sql"
|
||
|
||
async def main() -> None:
|
||
parser = argparse.ArgumentParser(description="初始化 mcp_auth 库并建表")
|
||
parser.add_argument("--host", default=os.getenv("MCP_AUTH_DB_HOST", "127.0.0.1"))
|
||
parser.add_argument("--port", type=int, default=int(os.getenv("MCP_AUTH_DB_PORT", "5432")))
|
||
parser.add_argument("--user", default=os.getenv("MCP_AUTH_DB_USER", "postgres"))
|
||
parser.add_argument("--password", default=os.getenv("MCP_AUTH_DB_PASSWORD", "digiwin"))
|
||
parser.add_argument("--db", default="mcp_auth")
|
||
args = parser.parse_args()
|
||
|
||
# 1. 连超库建 mcp_auth（IF NOT EXISTS）
|
||
sys_conn = await asyncpg.connect(
|
||
host=args.host, port=args.port, user=args.user, password=args.password, database="postgres"
|
||
)
|
||
exists = await sys_conn.fetchval("SELECT 1 FROM pg_database WHERE datname = $1", args.db)
|
||
if not exists:
|
||
await sys_conn.execute(f'CREATE DATABASE "{args.db}"')
|
||
print(f"[建库] {args.db} 已创建")
|
||
else:
|
||
print(f"[建库] {args.db} 已存在，跳过")
|
||
await sys_conn.close()
|
||
|
||
# 2. 连 mcp_auth 建表
|
||
conn = await asyncpg.connect(
|
||
host=args.host, port=args.port, user=args.user, password=args.password, database=args.db
|
||
)
|
||
init_sql = INIT_SQL_PATH.read_text(encoding="utf-8")
|
||
await conn.execute(init_sql)
|
||
print("[建表] mcp_token / mcp_token_log / admin_user / mcp_service 已就绪")
|
||
await conn.close()
|
||
|
||
print("\n[完成] 数据库初始化完毕，种子数据请通过管理后台操作")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
asyncio.run(main())
|