package db import ( "context" "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) } // ent 的 Schema.Create 只会"建新表",不会给已存在的表补列(本项目已知约束)。 // 新增字段统一在此以幂等 DDL 补齐(ADD COLUMN IF NOT EXISTS 可重复执行)。 applyColumnPatches(ctx) // 清理历史独立装箱表:装箱已合并到统一的出库主表 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 } // applyColumnPatches 幂等补齐已有表的新增列。 // 背景:ent 的 Schema.Create 对已存在表不会 ADD COLUMN,故新增字段必须在此登记, // 否则新库正常、存量库缺列导致运行时 500。 func applyColumnPatches(ctx context.Context) { if rawDB == nil { return } stmts := []string{ // P0-2 不合格品处置(问题记录 L135-145) `ALTER TABLE inspection_records ADD COLUMN IF NOT EXISTS disposal_type text NOT NULL DEFAULT 'NONE'`, `ALTER TABLE inspection_records ADD COLUMN IF NOT EXISTS disposal_remark text`, `ALTER TABLE inspection_records ADD COLUMN IF NOT EXISTS return_tracking_no text`, // D5 区域库位:货架号 / 第几层 / 位置号(手填) `ALTER TABLE zones ADD COLUMN IF NOT EXISTS shelf_no text`, `ALTER TABLE zones ADD COLUMN IF NOT EXISTS layer_no text`, `ALTER TABLE zones ADD COLUMN IF NOT EXISTS position_no text`, // D6 其他入库:归属(工单号/科技项目号/无) `ALTER TABLE inbound_orders ADD COLUMN IF NOT EXISTS ownership_type text NOT NULL DEFAULT 'none'`, `ALTER TABLE inbound_orders ADD COLUMN IF NOT EXISTS ownership_no text`, // 缺货预警:物料安全库存 `ALTER TABLE materials ADD COLUMN IF NOT EXISTS safety_stock bigint NOT NULL DEFAULT 0`, // 备料台账:目标工位(备料送到哪个装配工位) `ALTER TABLE order_material_ledgers ADD COLUMN IF NOT EXISTS target_station text`, } for _, s := range stmts { if _, err := rawDB.ExecContext(ctx, s); err != nil { slog.Warn("补齐列失败(可忽略): " + err.Error()) } } }