51 lines
1.0 KiB
Go
51 lines
1.0 KiB
Go
|
package db
|
||
|
|
|
||
|
|
import (
|
||
|
|
"database/sql"
|
||
|
|
"fmt"
|
||
|
|
"log/slog"
|
||
|
|
|
||
|
|
"bj_power_wms/ent"
|
||
|
|
|
||
|
|
"entgo.io/ent/dialect"
|
||
|
|
entsql "entgo.io/ent/dialect/sql"
|
||
|
|
_ "github.com/jackc/pgx/v5/stdlib"
|
||
|
|
)
|
||
|
|
|
||
|
|
// MustNewDB 创建 Ent 客户端并验证连接
|
||
|
|
func MustNewDB(c DatabaseConf) *ent.Client {
|
||
|
|
ec, err := NewDB(c)
|
||
|
|
if err != nil {
|
||
|
|
panic(fmt.Sprintf("数据库连接失败: %v", err))
|
||
|
|
}
|
||
|
|
return ec
|
||
|
|
}
|
||
|
|
|
||
|
|
func NewDB(c DatabaseConf) (*ent.Client, error) {
|
||
|
|
dsn := fmt.Sprintf("postgresql://%s:%s@%s:%d/%s",
|
||
|
|
c.User, c.Password, c.Host, c.Port, c.Dbname,
|
||
|
|
)
|
||
|
|
|
||
|
|
db, err := sql.Open("pgx", dsn)
|
||
|
|
if err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
|
||
|
|
if err := db.Ping(); err != nil {
|
||
|
|
return nil, fmt.Errorf("ping db: %w", err)
|
||
|
|
}
|
||
|
|
|
||
|
|
drv := entsql.OpenDB(dialect.Postgres, db)
|
||
|
|
return ent.NewClient(ent.Driver(drv)), nil
|
||
|
|
}
|
||
|
|
|
||
|
|
// AutoMigrate 启动时自动建表/迁移
|
||
|
|
func AutoMigrate(client *ent.Client) error {
|
||
|
|
ctx := contextBackdrop()
|
||
|
|
if err := client.Schema.Create(ctx); err != nil {
|
||
|
|
return fmt.Errorf("自动迁移失败: %w", err)
|
||
|
|
}
|
||
|
|
slog.Info("数据库迁移完成")
|
||
|
|
return nil
|
||
|
|
}
|