90 lines
2.7 KiB
Go
90 lines
2.7 KiB
Go
package db
|
||
|
||
import (
|
||
"context"
|
||
"fmt"
|
||
"os"
|
||
"strconv"
|
||
"time"
|
||
|
||
"github.com/jackc/pgx/v5/pgxpool"
|
||
)
|
||
|
||
// PoolProfile 描述应用池的连接预设（T20-B4：statement_timeout 分级）。
|
||
//
|
||
// 分级只通过池默认 + 显式 SET LOCAL 实现，业务 SQL 不允许散落裸 timeout：
|
||
// - default：statement_timeout=30s，常规 API 路径。
|
||
// - long：PI/索引/报表等长路径用 WithLongStatement 在事务内
|
||
// `SET LOCAL statement_timeout` 临时抬高，事务结束自动回落。
|
||
type PoolProfile struct {
|
||
// StatementTimeout 池默认语句超时；0 表示不设置（仅迁移会话使用）。
|
||
StatementTimeout time.Duration
|
||
// MaxConns 池上限；0 表示 pgx 默认（4×CPU）。
|
||
MaxConns int32
|
||
}
|
||
|
||
// DefaultProfile 常规业务池：30s 语句超时。
|
||
func DefaultProfile() PoolProfile {
|
||
return PoolProfile{
|
||
StatementTimeout: 30 * time.Second,
|
||
MaxConns: envInt32("APP_DB_MAX_CONNS", 8),
|
||
}
|
||
}
|
||
|
||
// MigrateProfile 迁移会话：不设语句超时（goose 链可能长）。
|
||
func MigrateProfile() PoolProfile {
|
||
return PoolProfile{StatementTimeout: 0, MaxConns: 2}
|
||
}
|
||
|
||
// parseWithProfile 构造 pool config 并注入 profile（供 OpenAppPool 与测试共用）。
|
||
func parseWithProfile(dsn string, profile PoolProfile) (*pgxpool.Config, error) {
|
||
cfg, err := pgxpool.ParseConfig(dsn)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("db pool: parse dsn: %w", err)
|
||
}
|
||
if profile.MaxConns > 0 {
|
||
cfg.MaxConns = profile.MaxConns
|
||
}
|
||
if profile.StatementTimeout > 0 {
|
||
if cfg.ConnConfig.RuntimeParams == nil {
|
||
cfg.ConnConfig.RuntimeParams = map[string]string{}
|
||
}
|
||
cfg.ConnConfig.RuntimeParams["statement_timeout"] = strconv.Itoa(
|
||
int(profile.StatementTimeout.Milliseconds()))
|
||
}
|
||
return cfg, nil
|
||
}
|
||
|
||
// OpenAppPool 以给定 profile 建池（T20-B4 统一入口）。
|
||
// dsn 透传 pgxpool；profile 控制 RuntimeParams 与池大小。
|
||
func OpenAppPool(ctx context.Context, dsn string, profile PoolProfile) (*pgxpool.Pool, error) {
|
||
cfg, err := parseWithProfile(dsn, profile)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
pool, err := pgxpool.NewWithConfig(ctx, cfg)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("db pool: connect: %w", err)
|
||
}
|
||
return pool, nil
|
||
}
|
||
|
||
// LongStatementTimeoutEnv 覆盖长路径语句超时（默认 15m，绑定 AGENT_TIMEOUT 量级）。
|
||
func LongStatementTimeout() time.Duration {
|
||
if v := os.Getenv("APP_DB_LONG_STATEMENT_TIMEOUT"); v != "" {
|
||
if d, err := time.ParseDuration(v); err == nil && d > 0 {
|
||
return d
|
||
}
|
||
}
|
||
return 15 * time.Minute
|
||
}
|
||
|
||
func envInt32(key string, fallback int32) int32 {
|
||
if v := os.Getenv(key); v != "" {
|
||
if n, err := strconv.ParseInt(v, 10, 32); err == nil && n > 0 {
|
||
return int32(n)
|
||
}
|
||
}
|
||
return fallback
|
||
}
|