初始化2

bj_power_wms/internal/db/db.go

@@ -0,0 +1,50 @@
+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
+}