228 lines
6.6 KiB
Go
228 lines
6.6 KiB
Go
// Package postgres 用 PostgreSQL 实现持久化存储接口（第一阶段：AuthKeyStore）。
|
||
//
|
||
// 查询代码由 sqlc 生成于 ./sqlcgen（见 telesrv/sqlc.yaml）；本包在其上实现 store 接口。
|
||
package postgres
|
||
|
||
import (
|
||
"context"
|
||
"errors"
|
||
"fmt"
|
||
"strings"
|
||
|
||
"github.com/golang-migrate/migrate/v4"
|
||
_ "github.com/golang-migrate/migrate/v4/database/pgx/v5" // 注册 pgx5:// migrate driver
|
||
"github.com/golang-migrate/migrate/v4/source/iofs"
|
||
"github.com/jackc/pgx/v5"
|
||
"github.com/jackc/pgx/v5/pgxpool"
|
||
|
||
"telesrv/deploy"
|
||
"telesrv/internal/store/postgres/sqlcgen"
|
||
)
|
||
|
||
const defaultMinConns = 16
|
||
|
||
const (
|
||
phoneIdentityPredecessorVersion = uint(181)
|
||
phoneIdentityMigrationVersion = uint(182)
|
||
// "phone182" as a signed PostgreSQL advisory-lock key. The lock spans the
|
||
// two-stage 0181 -> identity audit -> 0182 transition across new binaries.
|
||
phoneIdentityMigrationLockKey = int64(0x70686f6e65313832)
|
||
)
|
||
|
||
// MigrationStatus 是启动迁移后的 schema 状态。
|
||
type MigrationStatus struct {
|
||
Version uint
|
||
Dirty bool
|
||
Empty bool
|
||
}
|
||
|
||
// PoolOption 调整 pgxpool 连接池配置。
|
||
type PoolOption func(*pgxpool.Config)
|
||
|
||
// WithMaxConns 设置连接池最大连接数；<=0 时保持 pgx 默认。
|
||
// 同时把 MinConns 预热到 min(maxConns, 16)，降低 TDesktop 启动风暴下的冷连接尾延迟突刺。
|
||
func WithMaxConns(n int) PoolOption {
|
||
return func(cfg *pgxpool.Config) {
|
||
if n <= 0 {
|
||
return
|
||
}
|
||
cfg.MaxConns = int32(n)
|
||
minConns := int32(defaultMinConns)
|
||
if int32(n) < minConns {
|
||
minConns = int32(n)
|
||
}
|
||
cfg.MinConns = minConns
|
||
}
|
||
}
|
||
|
||
// WithMinConns 设置启动时预热的最小连接数；<=0 保持既有配置。
|
||
func WithMinConns(n int) PoolOption {
|
||
return func(cfg *pgxpool.Config) {
|
||
if n <= 0 {
|
||
return
|
||
}
|
||
minConns := int32(n)
|
||
if cfg.MaxConns > 0 && minConns > cfg.MaxConns {
|
||
minConns = cfg.MaxConns
|
||
}
|
||
cfg.MinConns = minConns
|
||
}
|
||
}
|
||
|
||
// Open 建立 pgxpool 连接池并 ping 验证。
|
||
func Open(ctx context.Context, dsn string, opts ...PoolOption) (*pgxpool.Pool, error) {
|
||
cfg, err := pgxpool.ParseConfig(dsn)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("pgxpool parse config: %w", err)
|
||
}
|
||
for _, opt := range opts {
|
||
if opt != nil {
|
||
opt(cfg)
|
||
}
|
||
}
|
||
cfg.ConnConfig.Tracer = queryStatsTracer{}
|
||
if err := installPostgresConnectionAdmission(ctx, cfg, dsn); err != nil {
|
||
return nil, fmt.Errorf("initialize PostgreSQL connection admission: %w", err)
|
||
}
|
||
pool, err := pgxpool.NewWithConfig(ctx, cfg)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("pgxpool new: %w", err)
|
||
}
|
||
if err := pool.Ping(ctx); err != nil {
|
||
pool.Close()
|
||
return nil, fmt.Errorf("pg ping: %w", err)
|
||
}
|
||
if err := warmMinConns(ctx, pool); err != nil {
|
||
pool.Close()
|
||
return nil, err
|
||
}
|
||
return pool, nil
|
||
}
|
||
|
||
func warmMinConns(ctx context.Context, pool *pgxpool.Pool) error {
|
||
target := pool.Config().MinConns
|
||
if target <= 0 {
|
||
return nil
|
||
}
|
||
conns := make([]*pgxpool.Conn, 0, target)
|
||
defer func() {
|
||
for _, conn := range conns {
|
||
conn.Release()
|
||
}
|
||
}()
|
||
for int32(len(conns)) < target {
|
||
conn, err := pool.Acquire(ctx)
|
||
if err != nil {
|
||
return fmt.Errorf("prewarm pg connection %d/%d: %w", len(conns)+1, target, err)
|
||
}
|
||
conns = append(conns, conn)
|
||
}
|
||
return nil
|
||
}
|
||
|
||
func withTx(ctx context.Context, db sqlcgen.DBTX, op string, fn func(pgx.Tx) error) error {
|
||
beginner, ok := db.(txBeginner)
|
||
if !ok {
|
||
return fmt.Errorf("%s: db does not support transactions", op)
|
||
}
|
||
tx, err := beginner.Begin(ctx)
|
||
if err != nil {
|
||
return fmt.Errorf("begin %s: %w", op, err)
|
||
}
|
||
committed := false
|
||
defer func() {
|
||
if !committed {
|
||
_ = tx.Rollback(ctx)
|
||
}
|
||
}()
|
||
if err := fn(tx); err != nil {
|
||
return err
|
||
}
|
||
if err := tx.Commit(ctx); err != nil {
|
||
return fmt.Errorf("commit %s: %w", op, err)
|
||
}
|
||
committed = true
|
||
return nil
|
||
}
|
||
|
||
// Migrate 用嵌入的迁移脚本将数据库迁移到最新版本。幂等：已最新时返回 nil。
|
||
func Migrate(dsn string) error {
|
||
_, err := MigrateAndStatus(dsn)
|
||
return err
|
||
}
|
||
|
||
// MigrateAndStatus 用嵌入迁移脚本迁移数据库，并返回迁移后的 schema 版本。
|
||
func MigrateAndStatus(dsn string) (MigrationStatus, error) {
|
||
ctx := context.Background()
|
||
lockConn, err := pgx.Connect(ctx, dsn)
|
||
if err != nil {
|
||
return MigrationStatus{}, fmt.Errorf("connect migration lock: %w", err)
|
||
}
|
||
defer lockConn.Close(ctx)
|
||
if _, err := lockConn.Exec(ctx, `SELECT pg_advisory_lock($1)`, phoneIdentityMigrationLockKey); err != nil {
|
||
return MigrationStatus{}, fmt.Errorf("lock migrations: %w", err)
|
||
}
|
||
defer func() {
|
||
_, _ = lockConn.Exec(context.Background(), `SELECT pg_advisory_unlock($1)`, phoneIdentityMigrationLockKey)
|
||
}()
|
||
|
||
src, err := iofs.New(deploy.Migrations, "migrations")
|
||
if err != nil {
|
||
return MigrationStatus{}, fmt.Errorf("iofs source: %w", err)
|
||
}
|
||
m, err := migrate.NewWithSourceInstance("iofs", src, toPgx5DSN(dsn))
|
||
if err != nil {
|
||
return MigrationStatus{}, fmt.Errorf("migrate new: %w", err)
|
||
}
|
||
defer m.Close()
|
||
status, err := migrationStatus(m)
|
||
if err != nil {
|
||
return MigrationStatus{}, err
|
||
}
|
||
if status.Dirty {
|
||
return MigrationStatus{}, fmt.Errorf("migrate version %d is dirty", status.Version)
|
||
}
|
||
if status.Empty || status.Version < phoneIdentityPredecessorVersion {
|
||
if err := m.Migrate(phoneIdentityPredecessorVersion); err != nil && !errors.Is(err, migrate.ErrNoChange) {
|
||
return MigrationStatus{}, fmt.Errorf("migrate to phone identity predecessor: %w", err)
|
||
}
|
||
status, err = migrationStatus(m)
|
||
if err != nil {
|
||
return MigrationStatus{}, err
|
||
}
|
||
if status.Dirty || status.Empty || status.Version != phoneIdentityPredecessorVersion {
|
||
return MigrationStatus{}, fmt.Errorf("phone identity predecessor status = %+v", status)
|
||
}
|
||
}
|
||
if status.Version < phoneIdentityMigrationVersion {
|
||
if err := canonicalizeStoredPhoneIdentities(ctx, lockConn); err != nil {
|
||
return MigrationStatus{}, fmt.Errorf("migrate phone identities: %w", err)
|
||
}
|
||
}
|
||
if err := m.Up(); err != nil && !errors.Is(err, migrate.ErrNoChange) {
|
||
return MigrationStatus{}, fmt.Errorf("migrate up: %w", err)
|
||
}
|
||
return migrationStatus(m)
|
||
}
|
||
|
||
func migrationStatus(m *migrate.Migrate) (MigrationStatus, error) {
|
||
version, dirty, err := m.Version()
|
||
if errors.Is(err, migrate.ErrNilVersion) {
|
||
return MigrationStatus{Empty: true}, nil
|
||
}
|
||
if err != nil {
|
||
return MigrationStatus{}, fmt.Errorf("migrate version: %w", err)
|
||
}
|
||
return MigrationStatus{Version: version, Dirty: dirty}, nil
|
||
}
|
||
|
||
// toPgx5DSN 把 pgxpool 用的 postgres:// DSN 转成 golang-migrate pgx5 driver 所需的 pgx5:// scheme。
|
||
func toPgx5DSN(dsn string) string {
|
||
if s, ok := strings.CutPrefix(dsn, "postgres://"); ok {
|
||
return "pgx5://" + s
|
||
}
|
||
if s, ok := strings.CutPrefix(dsn, "postgresql://"); ok {
|
||
return "pgx5://" + s
|
||
}
|
||
return dsn
|
||
}
|