lijiaoqiao/supply-api/internal/repository/db.go

package repository

import (
	"context"
	"fmt"
	"strings"
	"time"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
	"lijiaoqiao/supply-api/internal/config"
)

// DB 数据库连接池
type DB struct {
	Pool *pgxpool.Pool
}

// NewDB 创建数据库连接池
func NewDB(ctx context.Context, cfg config.DatabaseConfig) (*DB, error) {
	dsn := cfg.DSN()
	poolConfig, err := pgxpool.ParseConfig(dsn)
	if err != nil {
		// P2-05: 使用SafeDSN替代DSN，避免在错误信息中泄露密码
		return nil, fmt.Errorf("failed to parse database config for %s: %v", cfg.SafeDSN(), sanitizeErrorPassword(err, cfg.Password))
	}

	poolConfig.MaxConns = int32(cfg.MaxOpenConns)
	poolConfig.MinConns = int32(cfg.MaxIdleConns)
	poolConfig.MaxConnLifetime = cfg.ConnMaxLifetime
	poolConfig.MaxConnIdleTime = cfg.ConnMaxIdleTime
	poolConfig.HealthCheckPeriod = 30 * time.Second

	pool, err := pgxpool.NewWithConfig(ctx, poolConfig)
	if err != nil {
		// P2-05: 清理错误信息中的密码
		return nil, fmt.Errorf("failed to create connection pool for %s: %v", cfg.SafeDSN(), sanitizeErrorPassword(err, cfg.Password))
	}

	// 验证连接
	if err := pool.Ping(ctx); err != nil {
		pool.Close()
		return nil, fmt.Errorf("failed to ping database at %s:%d: %v", cfg.Host, cfg.Port, err)
	}

	return &DB{Pool: pool}, nil
}

// sanitizeErrorPassword 从错误信息中清理密码
// P2-05: pgxpool.ParseConfig的错误信息可能包含完整的DSN，需要清理
func sanitizeErrorPassword(err error, password string) error {
	if err == nil || password == "" {
		return err
	}
	// 将错误信息中的密码替换为***
	errStr := err.Error()
	safeErrStr := strings.ReplaceAll(errStr, password, "***")
	if safeErrStr != errStr {
		return fmt.Errorf("%s (password sanitized)", safeErrStr)
	}
	return err
}

// Close 关闭连接池
func (db *DB) Close() {
	if db.Pool != nil {
		db.Pool.Close()
	}
}

// HealthCheck 健康检查
func (db *DB) HealthCheck(ctx context.Context) error {
	return db.Pool.Ping(ctx)
}

// BeginTx 开始事务
func (db *DB) BeginTx(ctx context.Context) (Transaction, error) {
	tx, err := db.Pool.Begin(ctx)
	if err != nil {
		return nil, err
	}
	return &txWrapper{tx: tx}, nil
}

// Transaction 事务接口
type Transaction interface {
	Commit(ctx context.Context) error
	Rollback(ctx context.Context) error
}

type txWrapper struct {
	tx pgx.Tx
}

func (t *txWrapper) Commit(ctx context.Context) error {
	return t.tx.Commit(ctx)
}

func (t *txWrapper) Rollback(ctx context.Context) error {
	return t.tx.Rollback(ctx)
}