bj_power/bj_power_wms/internal/db/db.go

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`,
		// 备料台账：完成时间（最后一笔出库时间，问题记录 六.10）
		`ALTER TABLE order_material_ledgers ADD COLUMN IF NOT EXISTS completed_at bigint`,
		// 通用出库：出库类别/归属/目标工位（问题记录 二.2/二.3/二.4）
		`ALTER TABLE outbound_orders ADD COLUMN IF NOT EXISTS outbound_category text`,
		`ALTER TABLE outbound_orders ADD COLUMN IF NOT EXISTS ownership_type text`,
		`ALTER TABLE outbound_orders ADD COLUMN IF NOT EXISTS ownership_no text`,
		`ALTER TABLE outbound_orders ADD COLUMN IF NOT EXISTS target_station text`,
		// 库存明细：库位贯通 货架/层/位置（问题记录 一.4）
		`ALTER TABLE inventories ADD COLUMN IF NOT EXISTS shelf_no text`,
		`ALTER TABLE inventories ADD COLUMN IF NOT EXISTS layer_no text`,
		`ALTER TABLE inventories ADD COLUMN IF NOT EXISTS position_no text`,
		// 半成品/成品入库扩展（问题记录 L218/L222/L224/L230）
		`ALTER TABLE inventories ADD COLUMN IF NOT EXISTS ownership_type text`,
		`ALTER TABLE inventories ADD COLUMN IF NOT EXISTS ownership_no text`,
		`ALTER TABLE inventories ADD COLUMN IF NOT EXISTS product_status text`,
		`ALTER TABLE inventories ADD COLUMN IF NOT EXISTS related_standard text`,
		`ALTER TABLE inventories ADD COLUMN IF NOT EXISTS test_record_no text`,
		`ALTER TABLE inventories ADD COLUMN IF NOT EXISTS recorded_by text`,
		`ALTER TABLE inventories ADD COLUMN IF NOT EXISTS recorded_at bigint`,
		`ALTER TABLE inventories ADD COLUMN IF NOT EXISTS outbound_reason text`,
		// 检验记录扩展（问题记录 L107/L149/L150/L494：归属/相关标准/生产厂家/名称冗余）
		`ALTER TABLE inspection_records ADD COLUMN IF NOT EXISTS material_name text`,
		`ALTER TABLE inspection_records ADD COLUMN IF NOT EXISTS ownership_type text`,
		`ALTER TABLE inspection_records ADD COLUMN IF NOT EXISTS ownership_no text`,
		`ALTER TABLE inspection_records ADD COLUMN IF NOT EXISTS related_standard text`,
		`ALTER TABLE inspection_records ADD COLUMN IF NOT EXISTS manufacturer text`,
		// 入库单质量状态（问题记录 L129：入库记录增加 合格/不合格 列）
		`ALTER TABLE inbound_orders ADD COLUMN IF NOT EXISTS quality_status text`,
		// 入库单库位四级贯通（方案 E）：区域已有 zone_code，补货架/层/位置三列
		`ALTER TABLE inbound_orders ADD COLUMN IF NOT EXISTS shelf_no text`,
		`ALTER TABLE inbound_orders ADD COLUMN IF NOT EXISTS layer_no text`,
		`ALTER TABLE inbound_orders ADD COLUMN IF NOT EXISTS position_no text`,
		// 图号合并（问题记录 L18）：drawing_no 已并入 code，存量库丢弃旧列
		`ALTER TABLE materials DROP COLUMN IF EXISTS drawing_no`,
		// 批次号/SN 唯一约束（P2-1/P2-2）：并发导入防撞号。部分唯一索引只约束非空值——
		// 空串/NULL（电气件无批次、结构件无SN）不参与唯一性；batch_no 仅约束结构件(manage_mode=1)。
		// 注：存量若有重复数据会导致建索引失败（仅告警不阻断启动），上线前须先清理重复。
		`CREATE UNIQUE INDEX IF NOT EXISTS ux_inventories_batch_no ON inventories (batch_no) WHERE manage_mode = 1 AND batch_no IS NOT NULL AND batch_no <> ''`,
		`CREATE UNIQUE INDEX IF NOT EXISTS ux_inventories_sn_code ON inventories (sn_code) WHERE sn_code IS NOT NULL AND sn_code <> ''`,
		// U18 货位主数据 remodel：zone_code 从“单列唯一(大区时代)”降为“区域级(货位四级第一级)”，
		// 同一区域可登记多个货位；货位唯一性改由四级组合唯一索引保证，防止同一货位重复登记。
		// 先删 ent 旧生成的 zone_code 单列唯一索引，再建组合唯一索引（COALESCE 归一 NULL/空串，
		// 因 seed 的 Z01~Z04 三级明细为 NULL）。存量无重复四级组合（上线前已核实）。
		`DROP INDEX IF EXISTS zones_zone_code_key`,
		`CREATE UNIQUE INDEX IF NOT EXISTS ux_zones_location ON zones (zone_code, COALESCE(shelf_no, ''), COALESCE(layer_no, ''), COALESCE(position_no, ''))`,
		// U18 盘点四级贯通：盘点快照冻结货架/层/位置，盘点差异可精确定位到货位。
		`ALTER TABLE stocktake_items ADD COLUMN IF NOT EXISTS shelf_no text`,
		`ALTER TABLE stocktake_items ADD COLUMN IF NOT EXISTS layer_no text`,
		`ALTER TABLE stocktake_items ADD COLUMN IF NOT EXISTS position_no text`,
	}
	for _, s := range stmts {
		if _, err := rawDB.ExecContext(ctx, s); err != nil {
			slog.Warn("补齐列失败(可忽略): " + err.Error())
		}
	}
}