初始化2

bj_power_mes/internal/db/db.go

@@ -0,0 +1,57 @@
+package db
+
+import (
+	"database/sql"
+	"fmt"
+	"time"
+
+	"bj_power_mes/ent"
+
+	"entgo.io/ent/dialect"
+	entsql "entgo.io/ent/dialect/sql"
+	_ "github.com/jackc/pgx/v5/stdlib"
+)
+
+// DatabaseConf 数据库配置，DSN 由调用方从 yaml 读取
+type DatabaseConf struct {
+	Host     string `json:",default=127.0.0.1"`
+	Port     int    `json:",default=5432"`
+	User     string `json:",default=postgres"`
+	Password string `json:",default=postgres"`
+	Dbname   string `json:",default=bj_power_mes"`
+	MaxIdle  int    `json:",default=10"`
+	MaxOpen  int    `json:",default=20"`
+}
+
+// DSN 从配置拼装连接串（不在代码中硬编码口令）
+func (c DatabaseConf) DSN() string {
+	return fmt.Sprintf("host=%s port=%d user=%s password=%s dbname=%s sslmode=disable",
+		c.Host, c.Port, c.User, c.Password, c.Dbname)
+}
+
+// MustNewDB 打开 PostgreSQL 连接并包装为 ent.Client
+func MustNewDB(c DatabaseConf) (*ent.Client, *sql.DB) {
+	sqlDB, err := sql.Open("pgx", c.DSN())
+	if err != nil {
+		panic(fmt.Sprintf("open db failed: %v", err))
+	}
+	sqlDB.SetMaxIdleConns(c.MaxIdle)
+	sqlDB.SetMaxOpenConns(c.MaxOpen)
+	sqlDB.SetConnMaxLifetime(time.Hour)
+
+	drv := entsql.OpenDB(dialect.Postgres, sqlDB)
+	client := ent.NewClient(ent.Driver(drv))
+	return client, sqlDB
+}
+
+// NewDB 打开连接，失败返回 error（用于 migrate 等工具）
+func NewDB(c DatabaseConf) (*sql.DB, error) {
+	sqlDB, err := sql.Open("pgx", c.DSN())
+	if err != nil {
+		return nil, err
+	}
+	if err := sqlDB.Ping(); err != nil {
+		return nil, err
+	}
+	return sqlDB, nil
+}