57 lines
1.5 KiB
Go
57 lines
1.5 KiB
Go
|
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
|
||
|
|
}
|