合并分支

2023-08-04 17:11:10 +08:00
parent 12ec8d26bf
commit 020e76b901
100 changed files with 12692 additions and 2574 deletions
--- a/global/db.go
+++ b/global/db.go
@@ -1,60 +1,88 @@
 package global

 import (
-	"database/sql"
+	"context"
 	"fmt"
-	"time"

 	"electricity_bill_calc/config"
 	"electricity_bill_calc/logger"

-	"github.com/uptrace/bun"
-	"github.com/uptrace/bun/dialect/pgdialect"
-	"github.com/uptrace/bun/driver/pgdriver"
-	"go.uber.org/zap/zapcore"
+	"github.com/jackc/pgx/v5"
+	"github.com/jackc/pgx/v5/pgxpool"
+	"github.com/samber/lo"
+	"go.uber.org/zap"
 )

 var (
-	DB *bun.DB
+	DB *pgxpool.Pool
 )

 func SetupDatabaseConnection() error {
-	// connStr := fmt.Sprintf(
-	// 	"host=%s user=%s password=%s dbname=%s port=%d sslmode=disable TimeZone=Asia/Shanghai connect_timeout=0 tcp_user_timeout=180000",
-	// 	config.DatabaseSettings.Host,
-	// 	config.DatabaseSettings.User,
-	// 	config.DatabaseSettings.Pass,
-	// 	config.DatabaseSettings.DB,
-	// 	config.DatabaseSettings.Port,
-	// )
-	pgconn := pgdriver.NewConnector(
-		pgdriver.WithNetwork("tcp"),
-		pgdriver.WithAddr(fmt.Sprintf("%s:%d", config.DatabaseSettings.Host,
-			config.DatabaseSettings.Port)),
-		pgdriver.WithUser(config.DatabaseSettings.User),
-		pgdriver.WithInsecure(true),
-		pgdriver.WithPassword(config.DatabaseSettings.Pass),
-		pgdriver.WithDatabase(config.DatabaseSettings.DB),
-		pgdriver.WithDialTimeout(30*time.Second),
-		pgdriver.WithReadTimeout(3*time.Minute),
-		pgdriver.WithWriteTimeout(10*time.Minute),
+	connString := fmt.Sprintf(
+		"postgres://%s:%s@%s:%d/%s?sslmode=disable&connect_timeout=%d&application_name=%s&pool_max_conns=%d&pool_min_conns=%d&pool_max_conn_lifetime=%s&pool_max_conn_idle_time=%s&pool_health_check_period=%s",
+		config.DatabaseSettings.User,
+		config.DatabaseSettings.Pass,
+		config.DatabaseSettings.Host,
+		config.DatabaseSettings.Port,
+		config.DatabaseSettings.DB,
+		0,
+		"elec_service_go",
+		config.DatabaseSettings.MaxOpenConns,
+		config.DatabaseSettings.MaxIdleConns,
+		"60m",
+		"10m",
+		"10s",
 	)
-	sqldb := sql.OpenDB(pgconn)
-	DB = bun.NewDB(sqldb, pgdialect.New())
-	DB.AddQueryHook(logger.NewQueryHook(logger.QueryHookOptions{
-		LogSlow:         3 * time.Second,
-		Logger:          logger.Named("PG"),
-		QueryLevel:      zapcore.DebugLevel,
-		ErrorLevel:      zapcore.ErrorLevel,
-		SlowLevel:       zapcore.WarnLevel,
-		ErrorTemplate:   "{{.Operation}}[{{.Duration}}]: {{.Query}}: {{.Error}}",
-		MessageTemplate: "{{.Operation}}[{{.Duration}}]: {{.Query}}",
-	}))
-	DB.SetMaxIdleConns(config.DatabaseSettings.MaxIdleConns)
-	DB.SetMaxOpenConns(config.DatabaseSettings.MaxOpenConns)
-	DB.SetConnMaxIdleTime(10 * time.Minute)
-	DB.SetConnMaxLifetime(60 * time.Minute)
-	DB.Ping()
+	poolConfig, err := pgxpool.ParseConfig(connString)
+	if err != nil {
+		logger.Named("DB INIT").Error("数据库连接初始化失败。", zap.Error(err))
+		return err
+	}
+	poolConfig.ConnConfig.Tracer = QueryLogger{logger: logger.Named("PG")}
+	DB, _ = pgxpool.NewWithConfig(context.Background(), poolConfig)

 	return nil
 }
+
+type QueryLogger struct {
+	logger *zap.Logger
+}
+
+func (ql QueryLogger) TraceQueryStart(ctx context.Context, conn *pgx.Conn, data pgx.TraceQueryStartData) context.Context {
+	ql.logger.Info(fmt.Sprintf("将要执行查询: %s", data.SQL))
+	ql.logger.Info("查询参数", lo.Map(data.Args, func(elem any, index int) zap.Field {
+		return zap.Any(fmt.Sprintf("[Arg %d]: ", index), elem)
+	})...)
+	return ctx
+}
+
+func (ql QueryLogger) TraceQueryEnd(ctx context.Context, conn *pgx.Conn, data pgx.TraceQueryEndData) {
+	var logFunc func(string, ...zap.Field)
+	var templateString string
+	if data.Err != nil {
+		logFunc = ql.logger.Error
+		templateString = "命令 [%s] 执行失败。"
+	} else {
+		logFunc = ql.logger.Info
+		templateString = "命令 [%s] 执行成功。"
+	}
+	switch {
+	case data.CommandTag.Update():
+		fallthrough
+	case data.CommandTag.Delete():
+		fallthrough
+	case data.CommandTag.Insert():
+		logFunc(
+			fmt.Sprintf(templateString, data.CommandTag.String()),
+			zap.Error(data.Err),
+			zap.Any("affected", data.CommandTag.RowsAffected()))
+	case data.CommandTag.Select():
+		logFunc(
+			fmt.Sprintf(templateString, data.CommandTag.String()),
+			zap.Error(data.Err))
+	default:
+		logFunc(
+			fmt.Sprintf(templateString, data.CommandTag.String()),
+			zap.Error(data.Err))
+	}
+}