Compare commits

..

5 Commits

Author SHA1 Message Date
marco d46a5cd80f chore: alembic merge revision + PG 迁移指南文档
PostgreSQL 迁移 (#6) 之后两个分支引入了独立 heads (ef96beb47b1e +
d1e2f3a4b5c6), 必须合并成单一 head 才能 alembic upgrade head 顺利跑通。
- alembic/versions/f01db5d77dac_merge_pg_jsonb_and_feedback_heads.py:
  空 op 的 merge revision, 把两个 heads 收敛为 f01db5d77dac
- docs/postgres-migration.md: SQLite → PostgreSQL 迁移指南
  (背景/动机/操作步骤/无真实用户数据假设, 上线前补 pgloader 演练等)

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-05-30 18:08:51 +08:00
marco e69e244de7 Merge pull request 'feat(ad-reward): reward-status 加本轮 3 次冷却字段 (round_count + cooldown_until)' (#5) from feat/welfare-video-round-cooldown into main
Reviewed-on: #5
2026-05-29 15:45:51 +08:00
marco 2ebde935f9 Merge pull request 'feat: SQLite → PostgreSQL 迁移代码改动' (#6) from feat/postgres-migration-and-bootstrap-script into main
Reviewed-on: #6
2026-05-29 15:44:40 +08:00
chenshuobo b3d3fda744 feat: 新增 PG 初始化脚本 scripts/init_postgres.py
新机器一键拉起数据库, 替代手动跑 psql 建用户建库的 6 行命令.

脚本特性:
- 跨平台 Python (Windows/macOS/Linux 通吃, 不依赖 psql CLI 在 PATH 里)
- 直接用 psycopg 连 PG, 不 shell out
- 幂等: 用户已存在则重置密码, 库已存在则跳过, 重复跑不出错
- 业务密码自动 secrets.token_urlsafe(32) 强随机, 不让同事自己想
- 自动从 .env.example 复制出 .env 并写入 DATABASE_URL
- .env 里已有 DATABASE_URL 会询问是否覆盖 (不无脑覆盖)
- 跑完自动 alembic upgrade head, 一条命令到可用状态
- 全交互项支持环境变量提前指定 (PG_SUPER_PASS / APP_DB_PASS 等), CI 友好

同事新机器初始化流程从 6 步简化为 3 步:
  pip install -e .
  python scripts/init_postgres.py    # 唯一手动输入项: postgres 超级用户密码
  ./run.sh

CREATE/ALTER USER 注意点: PG 不支持参数化密码绑定, 用 sql.Literal 内联,
psycopg 自带的转义会处理引号和注入.

docs/数据库迁移.md 拆出 PG / SQLite 两套初始化流程, PG 走脚本,
SQLite 保留旧步骤 (本地临时回退用).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-05-29 15:21:04 +08:00
chenshuobo b76e5bd515 feat: SQLite → PostgreSQL 迁移代码改动
按 docs/数据库迁移.md 第 3 节扫雷清单落地代码层改动:

- pyproject.toml: 新增 psycopg[binary]>=3.1 依赖 (psycopg3, SQLAlchemy 2.0 时代默认, 不要再装 psycopg2)

- app/db/session.py: 非 SQLite 自动加连接池参数
  - pool_size=10, max_overflow=20, pool_recycle=3600
  - SQLite 单文件不池化, _is_sqlite 判断分支保留以便本地 dev 临时回退

- app/models/savings.py: dishes 列从 JSON 改为 JSONB
  - PG 上能用 GIN 索引和 jsonb 操作符
  - SQLite 上 JSON 实际是 TEXT, 切到 PG 后用 JSON 类型存的是 json 不是 jsonb

- alembic/versions/ef96beb47b1e_*.py: 新增迁移把 savings_record.dishes 从 json 转 jsonb
  - 用 op.get_bind().dialect.name 判断 PG 才执行 ALTER COLUMN
  - SQLite 上是 no-op (SQLite 没有 jsonb 类型)
  - 旧的 savings_shop_dishes.py 迁移保持 sa.JSON() 不动 (按文档"不改老迁移"原则)

本地切 PG 验证: 10 张表全部建出, dishes 字段为 jsonb 类型,
登录/me/美团 feed 三个真接口全 200, user 数据正确写入 PG.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-05-29 15:20:22 +08:00
8 changed files with 576 additions and 19 deletions
@@ -0,0 +1,36 @@
"""convert savings_record.dishes from json to jsonb
PG only. SQLite 上 JSON/JSONB 都是 TEXT, 此迁移是 no-op。
Revision ID: ef96beb47b1e
Revises: c8d9e0f1a2b3
Create Date: 2026-05-29 10:57:21.471774
"""
from typing import Sequence, Union
from alembic import op
from sqlalchemy.dialects.postgresql import JSONB
revision: str = 'ef96beb47b1e'
down_revision: Union[str, Sequence[str], None] = 'c8d9e0f1a2b3'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
if op.get_bind().dialect.name != "postgresql":
return
op.execute(
"ALTER TABLE savings_record "
"ALTER COLUMN dishes TYPE JSONB USING dishes::JSONB"
)
def downgrade() -> None:
if op.get_bind().dialect.name != "postgresql":
return
op.execute(
"ALTER TABLE savings_record "
"ALTER COLUMN dishes TYPE JSON USING dishes::JSON"
)
@@ -0,0 +1,26 @@
"""merge pg_jsonb and feedback heads
Revision ID: f01db5d77dac
Revises: ef96beb47b1e, d1e2f3a4b5c6
Create Date: 2026-05-29 16:09:58.498935
"""
from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa
# revision identifiers, used by Alembic.
revision: str = 'f01db5d77dac'
down_revision: Union[str, Sequence[str], None] = ('ef96beb47b1e', 'd1e2f3a4b5c6')
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
pass
def downgrade() -> None:
pass
+14 -8
View File
@@ -27,19 +27,25 @@ def _ensure_sqlite_dir(url: str) -> None:
_ensure_sqlite_dir(settings.DATABASE_URL)
_is_sqlite = settings.DATABASE_URL.startswith("sqlite")
# SQLite 跨线程访问要 check_same_thread=False;PG/MySQL 不需要这个参数
_connect_args: dict = {}
if settings.DATABASE_URL.startswith("sqlite"):
if _is_sqlite:
_connect_args["check_same_thread"] = False
engine = create_engine(
settings.DATABASE_URL,
connect_args=_connect_args,
_engine_kwargs: dict = {
"connect_args": _connect_args,
# echo 在 dev 下打 SQL,生产关掉
echo=settings.APP_DEBUG and not settings.is_prod,
future=True,
pool_pre_ping=True,
)
"echo": settings.APP_DEBUG and not settings.is_prod,
"future": True,
"pool_pre_ping": True,
}
# SQLite 用单文件不需要池;PG/MySQL 必须显式池化 + recycle 防 idle 断连
if not _is_sqlite:
_engine_kwargs.update(pool_size=10, max_overflow=20, pool_recycle=3600)
engine = create_engine(settings.DATABASE_URL, **_engine_kwargs)
SessionLocal = sessionmaker(bind=engine, autocommit=False, autoflush=False, expire_on_commit=False)
+4 -3
View File
@@ -8,7 +8,8 @@ from __future__ import annotations
from datetime import datetime
from sqlalchemy import JSON, DateTime, ForeignKey, Integer, String, func
from sqlalchemy import DateTime, ForeignKey, Integer, String, func
from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy.orm import Mapped, mapped_column
from app.db.base import Base
@@ -28,8 +29,8 @@ class SavingsRecord(Base):
title: Mapped[str | None] = mapped_column(String(128), nullable=True)
# 店铺名(订单明细卡标题,如「窑鸡王(王府井店)」)
shop_name: Mapped[str | None] = mapped_column(String(128), nullable=True)
# 菜品名列表(JSON),前 2 道直接展示,其余收进「还有 N 道菜」展开。SQLite 存为 TEXT。
dishes: Mapped[list[str]] = mapped_column(JSON, nullable=False, default=list)
# 菜品名列表(JSONB),PG 上可建 GIN 索引,前 2 道直接展示,其余收进「还有 N 道菜」展开。
dishes: Mapped[list[str]] = mapped_column(JSONB, nullable=False, default=list)
# 来源:demo(演示) / compare(真实比价上报)
source: Mapped[str] = mapped_column(String(16), nullable=False, default="compare")
+279
View File
@@ -0,0 +1,279 @@
# SQLite → PostgreSQL 迁移指南
> 上线前把数据库从 SQLite 切到 PostgreSQL。本服务自始就为这一步预留了路径(`db/session.py` 的 dialect 特判 / `alembic/env.py` 的 batch mode 开关 / SQLAlchemy 2.0 + Alembic),代码改动很小。
>
> **背景假设**:当前生产**无真实用户数据**(MVP 阶段),所以整个迁移本质是"切引擎",不涉及数据搬迁。如果将来已有真实数据,本文档不适用,需补充 `pgloader` 演练 + 钱表金额逐行核对 + 停服窗口 + 回滚预案。
>
> 最后更新:2026-05-29
---
## 0. 为什么切
SQLite 三个硬伤决定它不能上线:
- **并发写锁库**:整个文件一把锁,多请求并发写会排队,极易超时。当前 uvicorn `--workers 1` 勉强能跑,扩 worker 立刻撞墙。
- **类型宽松**:整数列能塞字符串、`CHECK` 约束当装饰品。钱表(`wallet` / `withdraw_order` / `cash_transaction`)用 SQLite 等于裸奔。
- **无 JSONB / 无并发索引 / 无分区**:PG 的核心能力一个都享受不到。`savings.dishes` 现在用 SQLite 的 `JSON`(实际存 TEXT),无索引、无操作符,纯展示用。
PG 默认上 16 版(工具链最齐),驱动用 **psycopg3**(SQLAlchemy 2.0 时代默认,不要再装 psycopg2)。
---
## 1. 本地起 PG + 跑通空库(半天)
### 1.1 装 PG
macOS:
```bash
brew install postgresql@16
brew services start postgresql@16
```
Linux(Ubuntu/Debian):
```bash
sudo apt install postgresql-16
sudo systemctl enable --now postgresql
```
### 1.2 建库 + 用户 + 授权
```bash
psql postgres <<'EOF'
CREATE USER shaguabijia_app WITH PASSWORD 'change-me-strong-random';
CREATE DATABASE shaguabijia OWNER shaguabijia_app ENCODING 'UTF8';
GRANT ALL PRIVILEGES ON DATABASE shaguabijia TO shaguabijia_app;
EOF
```
密码用 `python -c "import secrets; print(secrets.token_urlsafe(32))"` 生成,**不要复用 JWT_SECRET_KEY**。
### 1.3 装驱动
```bash
pip install "psycopg[binary]>=3.1"
```
`psycopg[binary]` 是预编译版,免装 libpq 头文件。生产同款,无需源码编译版。
### 1.4 改 `.env`
```ini
DATABASE_URL=postgresql+psycopg://shaguabijia_app:<password>@localhost:5432/shaguabijia
```
> ⚠️ URL scheme 必须是 `postgresql+psycopg://`(显式声明 psycopg3),不能写成 `postgresql://`——SQLAlchemy 默认会去找 psycopg2,装的不一致就报 ModuleNotFoundError。
### 1.5 建表
```bash
alembic upgrade head
```
### 1.6 验证
```bash
psql -U shaguabijia_app -d shaguabijia -c "\dt"
```
应该列出全部 9 张表 + `alembic_version`:
```
user, coin_account, coin_transaction, cash_transaction,
withdraw_order, signin_record, task_claim_record,
savings_record, ad_reward_record, alembic_version
```
再手动 INSERT 一行 user,确认 `created_at` 默认值正常落:
```sql
INSERT INTO "user" (phone, register_channel, status) VALUES ('13800138000', 'sms', 'active');
SELECT id, phone, created_at, last_login_at FROM "user";
```
`created_at` 应该是当前时间戳带时区,不是 NULL。`user` 表名是 PG 关键字,**必须用双引号**包裹。
---
## 2. 测试套件全绿(1 小时)
### 2.1 测试连 PG
测试默认连 SQLite。临时让 `tests/conftest.py``DATABASE_URL` 也指到本地 PG(可以建个 `shaguabijia_test` 库专门跑测试,免得污染 dev 库)。
```bash
createdb -O shaguabijia_app shaguabijia_test
DATABASE_URL=postgresql+psycopg://shaguabijia_app:xxx@localhost:5432/shaguabijia_test pytest
```
### 2.2 处理失败用例
**所有红色测试都不能放过**——SQLite 类型/约束都很宽松,PG 严格起来暴露的就是真 bug。常见问题:
- **字符串和整数比较**:SQLite 允许 `WHERE phone = 13800138000`(自动转字符串),PG 直接报错。改成 `:phone` 参数 + 类型严格
- **timezone 处理**:SQLite 存 naive datetime,PG 的 `TIMESTAMPTZ` 必须带时区。检查所有 `datetime.utcnow()`,改成 `datetime.now(timezone.utc)`
- **隐式 commit**:SQLite 某些情况下隐式 commit,PG 严格事务边界。如果某用例报 `current transaction is aborted`,说明业务代码缺 commit/rollback
### 2.3 验证
`pytest` 全绿,**且本地起 `./run.sh` 后用 Apifox 调通登录 / 美团 feed / 领券透传三个真接口**。
---
## 3. 代码扫雷(1 天)
### 3.1 必改清单
| 文件 / 位置 | 改动 | 原因 |
|---|---|---|
| `app/models/savings.py` | `from sqlalchemy import JSON``from sqlalchemy.dialects.postgresql import JSONB`;`mapped_column(JSON, ...)``mapped_column(JSONB, ...)` | JSONB 在 PG 上有 GIN 索引和操作符支持,是 PG 核心优势之一。SQLite 上 `JSON` 实际是 TEXT,切到 PG 后用 `JSON` 类型存的是 `json` 不是 `jsonb`,查询/索引能力差一大截 |
| `app/db/session.py``create_engine` 调用 | 增加连接池参数:`pool_size=10, max_overflow=20, pool_recycle=3600` | SQLite 单文件不需要池,PG 必须显式池化。`pool_recycle=3600` 防 PG 主动断开的 idle 连接 |
| `alembic/versions/` 下旧迁移里的 `sa.text('(CURRENT_TIMESTAMP)')` | **不要改** | PG 能接受这个语法。改了反而搞乱迁移历史 + 让 alembic 觉得 schema drift。新写的迁移统一用 `sa.func.now()` |
### 3.2 保留的 dialect 特判
下面这些**保留原样**,它们的 SQLite 分支在 PG 下自动绕过、不影响新部署,但本地开发想用 SQLite 临时跑还能用:
- `app/db/session.py` 里的 `_ensure_sqlite_dir``check_same_thread` 特判
- `alembic/env.py` 里的 `render_as_batch` 特判
### 3.3 顺便清掉
`app/main.py` 启动日志里如果打 `DATABASE_URL`,确认密码不会进日志(只打 scheme + host,不打完整 URL)。
### 3.4 验证
本地 PG 起 `./run.sh`,跑一遍下面的烟测脚本:
```bash
# 登录
curl -X POST http://localhost:8770/api/v1/auth/sms/send -d '{"phone":"13800138000"}'
curl -X POST http://localhost:8770/api/v1/auth/sms/login -d '{"phone":"13800138000","code":"123456"}'
# 拿到 access_token 后调 me
curl -H "Authorization: Bearer <token>" http://localhost:8770/api/v1/auth/me
# 美团 feed
curl -X POST http://localhost:8770/api/v1/meituan/feed -d '{"longitude":116.4,"latitude":39.9,"page":0}'
```
三条都 200 即通过。`psql``user` 表多了一条 `13800138000`
---
## 4. 生产切换(半小时,无需停服窗口)
**前提**:生产无真实数据。如已有数据,本节不适用。
### 4.1 装 PG
```bash
ssh server
sudo apt install postgresql-16
sudo systemctl enable --now postgresql
```
### 4.2 建库
```bash
sudo -u postgres psql <<'EOF'
CREATE USER shaguabijia_app WITH PASSWORD '<strong-random-password>';
CREATE DATABASE shaguabijia OWNER shaguabijia_app ENCODING 'UTF8';
GRANT ALL PRIVILEGES ON DATABASE shaguabijia TO shaguabijia_app;
EOF
```
生产密码**单独生成**,不要复用本地的。
### 4.3 装驱动到生产 venv
```bash
cd /opt/shaguabijia-app-server
.venv/bin/pip install "psycopg[binary]>=3.1"
```
### 4.4 改生产 `.env`
```ini
DATABASE_URL=postgresql+psycopg://shaguabijia_app:<password>@localhost:5432/shaguabijia
```
### 4.5 建表 + 重启
```bash
.venv/bin/alembic upgrade head
sudo systemctl restart shaguabijia-app-server
sudo journalctl -u shaguabijia-app-server -f
```
启动日志里看到 `Application startup complete` + 没有 5xx 即通过。
### 4.6 验证
- 用真账号(手机号)走一遍极光登录或短信登录
-`GET /api/v1/auth/me`,确认返回当前用户
-`POST /api/v1/meituan/feed`,确认券列表正常
- `psql -U shaguabijia_app -d shaguabijia -c 'SELECT id, phone, created_at FROM "user"'` 看到刚登录的用户
### 4.7 老 SQLite 文件归档
```bash
cd /opt/shaguabijia-app-server
mv data/app.db data/app.db.legacy-$(date +%Y%m%d)
```
不要删——万一要回查测试期残留数据(开发期可能塞了一些手工记录),保留几个月。
---
## 5. 兜底:切换失败如何回滚
万一改完起不来:
```bash
# 1. .env 改回 SQLite
sed -i 's|^DATABASE_URL=postgresql.*|DATABASE_URL=sqlite:///./data/app.db|' /opt/shaguabijia-app-server/.env
# 2. 把归档的 sqlite 文件还原(如果已经改名)
mv data/app.db.legacy-* data/app.db
# 3. 重启
sudo systemctl restart shaguabijia-app-server
```
由于 SQLite 文件没动过,无数据丢失,损失只是切换过程中的几分钟 5xx。
---
## 6. 顺便建议:同时装上 Redis
PG 上线后,Redis 是下一个明确要补的基础设施(已在 [待办与技术债.md](./待办与技术债.md) 里列了用途)。**建议本次维护窗口顺手把 Redis 也装上**,不开始用、只装服务,免得下次还要单独申请部署:
```bash
sudo apt install redis-server
sudo systemctl enable --now redis-server
# 改 /etc/redis/redis.conf:bind 127.0.0.1(只本机访问)+ requirepass <strong-password>
sudo systemctl restart redis-server
```
下面三个用途等接的时候再写代码,**本次只装服务、不接代码**:
| 用途 | 当前临时方案 | 何时必须接 Redis |
|---|---|---|
| 短信验证码冷却 | 进程内存 dict(`integrations/sms.py`) | `SMS_MOCK=false` 上线时 / 扩 worker 时 |
| pricebot trace_id session | 进程内存 dict(在 pricebot-backend) | pricebot 扩 worker 时 |
| JWT 黑名单(logout 真失效) | 当前 logout 占位无失效 | P1 鉴权改造时 |
---
## 7. 后续优化(不阻塞上线,一周内补)
- **每天凌晨 `pg_dump` 备份**:用 systemd timer 或 cron,保留 7 天。备份文件 `pg_dump -Fc shaguabijia > /backup/shaguabijia-$(date +%Y%m%d).dump`
- **金钱表加 CHECK 约束**:新迁移给 `coin_account.balance``withdraw_order.amount_yuan_x100``CHECK (... >= 0)`。PG 严格执行,SQLite 形同虚设
- **timezone 统一**:生产 PG 配 `timezone='Asia/Shanghai'` 或维持 `UTC`,**团队定一个**写进本文档,业务代码统一处理。当前模型混用 `func.now()`(server-side)和 `_utcnow()`(`models/user.py``last_login_at`,Python-side),后者依赖应用层时区,迁完 PG 后建议统一改 `func.now()`
- **慢查询观察**:`shared_preload_libraries = 'pg_stat_statements'` 打开扩展,每周看一次 Top 10 慢查询,提前优化
---
## 8. 不要做的事
- ❌ 不要搞主从复制 / 读写分离——量级远没到,先用单实例
- ❌ 不要装 PgBouncer / 连接池中间件——SQLAlchemy 自带池够用
- ❌ 不要 squash 旧迁移文件——9 个增量保持原样,让 alembic 一次跑过空库
- ❌ 不要在生产 PG 上裸跑测试套件——用单独的 `shaguabijia_test`
+25 -8
View File
@@ -15,20 +15,37 @@ alembic upgrade head
---
## 1. 首次初始化(clone 后)
### 1.1 用 PostgreSQL(推荐,与生产一致)
前置:本机已装 PostgreSQL 16 + 知道 postgres 超级用户密码。
```bash
# (1) 装依赖(在你的虚拟环境里)
# (1) 装依赖(在你的虚拟环境里, 含 psycopg)
pip install -e .
# (2) 配 .env(从模板复制,至少填 JWT_SECRET_KEY)
cp .env.example .env
# (2) 一键建用户 + 建库 + 写 .env + 跑迁移
python scripts/init_postgres.py
# 按提示输入 PG 超级用户密码即可。脚本会:
# - 自动建业务用户 shaguabijia_app + 业务库 shaguabijia
# - 自动生成业务用户强密码,写入 .env 的 DATABASE_URL
# - 跑完 alembic upgrade head
# (3) SQLite 默认库需要 data/ 目录存在
# (3) 启动
./run.sh
```
可用环境变量提前指定(免交互):
```bash
PG_SUPER_PASS=xxx APP_DB_PASS=yyy python scripts/init_postgres.py
```
脚本幂等,重复跑会重置业务用户密码、跳过已存在的库。
### 1.2 用 SQLite(本地开发临时用,不接生产链路时)
```bash
pip install -e .
cp .env.example .env # 不动 DATABASE_URL, 默认 sqlite:///./data/app.db
mkdir -p data
# (4) 跑迁移建表 —— 关键
alembic upgrade head
# (5) 启动(run.sh 会自动重跑 3+4,所以平时直接 ./run.sh 也行)
./run.sh
```
> 也可以直接 `bash scripts/migrate.sh` 只做迁移、不启服务(部署/CI 用)。
+3
View File
@@ -17,6 +17,9 @@ dependencies = [
"sqlalchemy>=2.0.35",
"alembic>=1.13.3",
# PostgreSQL 驱动 (psycopg3, SQLAlchemy 2.0 时代默认, 不要再装 psycopg2)
"psycopg[binary]>=3.1",
# JWT 签名 / 校验
"pyjwt[crypto]>=2.9.0",
+189
View File
@@ -0,0 +1,189 @@
"""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())