130 lines
4.1 KiB
Python
130 lines
4.1 KiB
Python
"""PostgreSQL 测试库准备（一次性）。
|
||
|
||
做三件事：
|
||
1. 探测本机 PostgreSQL 连接（默认 127.0.0.1:5433，用户 postgres）；
|
||
2. 创建 `dealerhub`（开发）与 `dealerhub_test`（测试）两个库（幂等）；
|
||
3. 应用迁移到开发库，使 `DJANGO_SETTINGS_MODULE=config.settings.dev` 能直接连 PG 跑。
|
||
|
||
用法：
|
||
python scripts/setup_pg.py # 用默认凭据探测
|
||
PG_PORT=5432 PG_USER=me PG_PASSWORD=xx python scripts/setup_pg.py
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import os
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
||
|
||
PG_HOST = os.environ.get("PG_HOST", "127.0.0.1")
|
||
PG_PORT = int(os.environ.get("PG_PORT", "5433"))
|
||
PG_USER = os.environ.get("PG_USER", "postgres")
|
||
PG_PASSWORD = os.environ.get("PG_PASSWORD", "postgres")
|
||
|
||
DEV_DB = os.environ.get("PG_DEV_DB", "dealerhub")
|
||
TEST_DB = os.environ.get("PG_TEST_DB", "dealerhub_test")
|
||
|
||
# 常见本地 PG 密码（只用于"探测"，不做暴力破解）
|
||
PASSWORD_CANDIDATES = [PG_PASSWORD, "", "postgres", "123456", "root", "pg123456", "admin"]
|
||
|
||
|
||
def probe_connection() -> str | None:
|
||
"""返回可用密码；失败返回 None。"""
|
||
try:
|
||
import psycopg
|
||
except ImportError:
|
||
print("✗ 未安装 psycopg：pip install 'psycopg[binary]'")
|
||
return None
|
||
|
||
seen = []
|
||
for pw in PASSWORD_CANDIDATES:
|
||
if pw in seen:
|
||
continue
|
||
seen.append(pw)
|
||
try:
|
||
conn = psycopg.connect(
|
||
host=PG_HOST, port=PG_PORT, user=PG_USER, password=pw,
|
||
dbname="postgres", connect_timeout=3,
|
||
)
|
||
conn.close()
|
||
return pw
|
||
except Exception as exc:
|
||
print(f" · {PG_USER}/{pw or '(空)'} → {str(exc)[:70]}")
|
||
return None
|
||
|
||
|
||
def ensure_databases(password: str) -> bool:
|
||
import psycopg
|
||
|
||
ok = True
|
||
conn = psycopg.connect(
|
||
host=PG_HOST, port=PG_PORT, user=PG_USER, password=password,
|
||
dbname="postgres", autocommit=True,
|
||
)
|
||
cur = conn.cursor()
|
||
for db in (DEV_DB, TEST_DB):
|
||
cur.execute("SELECT 1 FROM pg_database WHERE datname = %s", (db,))
|
||
if cur.fetchone():
|
||
print(f" · {db} 已存在")
|
||
else:
|
||
cur.execute(f'CREATE DATABASE "{db}"')
|
||
print(f" ✓ 创建 {db}")
|
||
cur.close()
|
||
conn.close()
|
||
return ok
|
||
|
||
|
||
def apply_migrations_to_dev(password: str) -> bool:
|
||
"""把迁移应用到开发库（让 dev 设置能直接用 PG）。"""
|
||
os.environ.update({
|
||
"DATABASE_URL": f"postgres://{PG_USER}:{password}@{PG_HOST}:{PG_PORT}/{DEV_DB}",
|
||
"DJANGO_SETTINGS_MODULE": "config.settings.dev",
|
||
})
|
||
import django
|
||
|
||
django.setup()
|
||
from django.core.management import call_command
|
||
|
||
call_command("migrate", "--noinput", verbosity=0)
|
||
return True
|
||
|
||
|
||
def main() -> int:
|
||
print(f"探测 PostgreSQL · {PG_HOST}:{PG_PORT} · user={PG_USER}")
|
||
pw = probe_connection()
|
||
if pw is None:
|
||
print("\n✗ 无法连接。请确认：")
|
||
print(" 1. PostgreSQL 服务已启动")
|
||
print(" 2. 端口正确（PG_PORT=5433 是本机 PG 16 的常见端口）")
|
||
print(" 3. 凭据正确（PG_USER / PG_PASSWORD）")
|
||
return 1
|
||
print(f" ✓ 连接成功（密码：{pw or '(空)'}）")
|
||
|
||
print("\n创建数据库")
|
||
ensure_databases(pw)
|
||
|
||
print("\n应用迁移到开发库")
|
||
apply_migrations_to_dev(pw)
|
||
print(f" ✓ {DEV_DB} 迁移完成")
|
||
|
||
print(f"""
|
||
准备完成。运行方式：
|
||
|
||
# 开发/手工验证（连 PG）
|
||
export DATABASE_URL="postgres://{PG_USER}:{pw}@{PG_HOST}:{PG_PORT}/{DEV_DB}"
|
||
DJANGO_SETTINGS_MODULE=config.settings.dev granian config.asgi:application --interface asgi --port 8000
|
||
|
||
# 全量测试（含并发，PG 行锁语义）
|
||
DJANGO_SETTINGS_MODULE=config.settings.pgtest python -m pytest tests/ -p no:cacheprovider
|
||
|
||
# 只跑并发用例
|
||
DJANGO_SETTINGS_MODULE=config.settings.pgtest python -m pytest tests/test_concurrency.py -p no:cacheprovider -m postgres
|
||
""")
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
sys.exit(main())
|