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 }