bj_power/bj_power_wms/internal/db/db.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"
)

// rawDB 持有底层 *sql.DB，供迁移期清理孤儿表（如已下线的 package_boxes）使用
var rawDB *sql.DB

// 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?sslmode=disable",
		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)
	}

	rawDB = db
	drv := entsql.OpenDB(dialect.Postgres, db)
	return ent.NewClient(ent.Driver(drv)), nil
}

// EnsureDB 若目标数据库不存在则自动创建（先连 postgres 维护库做检查/创建）
func EnsureDB(c DatabaseConf) error {
	maintenanceDSN := fmt.Sprintf("postgresql://%s:%s@%s:%d/%s?sslmode=disable",
		c.User, c.Password, c.Host, c.Port, "postgres")
	sqlDB, err := sql.Open("pgx", maintenanceDSN)
	if err != nil {
		return err
	}
	defer sqlDB.Close()

	var exists bool
	if err := sqlDB.QueryRow(
		`SELECT EXISTS(SELECT 1 FROM pg_database WHERE datname = $1)`, c.Dbname,
	).Scan(&exists); err != nil {
		return err
	}

	if exists {
		return nil
	}

	if _, err := sqlDB.Exec(`CREATE DATABASE "` + c.Dbname + `"`); err != nil {
		return err
	}
	slog.Info("已自动创建数据库 " + c.Dbname)
	return nil
}

// AutoMigrate 启动时自动建表/迁移
func AutoMigrate(client *ent.Client) error {
	ctx := contextBackdrop()
	if err := client.Schema.Create(ctx); err != nil {
		return fmt.Errorf("自动迁移失败: %w", err)
	}
	// 清理历史独立装箱表：装箱已合并到统一的出库主表 OutboundOrder，
	// 原 package_boxes 表成为孤儿表，此处幂等删除（DROP TABLE IF EXISTS 可重复执行）。
	if rawDB != nil {
		if _, err := rawDB.ExecContext(ctx, "DROP TABLE IF EXISTS package_boxes"); err != nil {
			slog.Warn("清理旧装箱表 package_boxes 失败(可忽略): " + err.Error())
		} else {
			slog.Info("已清理旧装箱表 package_boxes")
		}
	}
	slog.Info("数据库迁移完成")
	return nil
}