bj_power/bj_power_mes/internal/db/db.go

package db

import (
	"context"
	"database/sql"
	"fmt"
	"time"

	"bj_power_mes/ent"

	"entgo.io/ent/dialect"
	entsql "entgo.io/ent/dialect/sql"
	_ "github.com/jackc/pgx/v5/stdlib"
)

// DatabaseConf 数据库配置，DSN 由调用方从 yaml 读取
type DatabaseConf struct {
	Host     string `json:",default=127.0.0.1"`
	Port     int    `json:",default=5432"`
	User     string `json:",default=postgres"`
	Password string `json:",default=postgres"`
	Dbname   string `json:",default=bj_power_mes"`
	MaxIdle  int    `json:",default=10"`
	MaxOpen  int    `json:",default=20"`
}

// DSN 从配置拼装连接串（不在代码中硬编码口令）
func (c DatabaseConf) DSN() string {
	return fmt.Sprintf("host=%s port=%d user=%s password=%s dbname=%s sslmode=disable",
		c.Host, c.Port, c.User, c.Password, c.Dbname)
}

// MustNewDB 打开 PostgreSQL 连接并包装为 ent.Client
func MustNewDB(c DatabaseConf) (*ent.Client, *sql.DB) {
	sqlDB, err := sql.Open("pgx", c.DSN())
	if err != nil {
		panic(fmt.Sprintf("open db failed: %v", err))
	}
	sqlDB.SetMaxIdleConns(c.MaxIdle)
	sqlDB.SetMaxOpenConns(c.MaxOpen)
	sqlDB.SetConnMaxLifetime(time.Hour)

	drv := entsql.OpenDB(dialect.Postgres, sqlDB)
	client := ent.NewClient(ent.Driver(drv))
	return client, sqlDB
}

// NewDB 打开连接，失败返回 error（用于 migrate 等工具）
func NewDB(c DatabaseConf) (*sql.DB, error) {
	sqlDB, err := sql.Open("pgx", c.DSN())
	if err != nil {
		return nil, err
	}
	if err := sqlDB.Ping(); err != nil {
		return nil, err
	}
	return sqlDB, nil
}

// EnsureDB 若目标数据库不存在则自动创建。
// 先连接 PostgreSQL 内置维护库 maintenance db（默认 postgres），检查目标库，
// 不存在则执行 CREATE DATABASE，便于重置/建表命令在空环境一键运行。
func EnsureDB(c DatabaseConf) error {
	maintenanceDSN := fmt.Sprintf("host=%s port=%d user=%s password=%s dbname=postgres sslmode=disable",
		c.Host, c.Port, c.User, c.Password)
	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
	}
	return nil
}

// schemaPatchSQL 幂等列补丁：ent Schema.Create 只建缺失的新表，绝不会给既有表加列；
// 因此对"已有表新增列"的场景必须走显式 ALTER。这里统一维护补丁清单，migrate/seed/reset-all 三入口共用。
var schemaPatchSQL = []string{
	// BOM 料增加"装配工序"维度（0=不参与工位绑定校验；1~12=在第几道工序装配）
	// 注意：BomItem schema 表名注解为 work_order_bom（非默认名）
	`ALTER TABLE work_order_bom ADD COLUMN IF NOT EXISTS process_code integer NOT NULL DEFAULT 0`,

	// 第三批：工程编号三级归属（合同号 → 工程编号 → 产品序号）
	`ALTER TABLE work_order ADD COLUMN IF NOT EXISTS contract_no varchar(64) NOT NULL DEFAULT ''`,
	`ALTER TABLE work_order ADD COLUMN IF NOT EXISTS project_no varchar(64) NOT NULL DEFAULT ''`,
	`ALTER TABLE work_order ADD COLUMN IF NOT EXISTS product_serial varchar(64) NOT NULL DEFAULT ''`,

	// P0-4：过程巡检按步骤提交（步骤维度字段）
	`ALTER TABLE inspection_record ADD COLUMN IF NOT EXISTS process_code integer NOT NULL DEFAULT 0`,
	`ALTER TABLE inspection_record ADD COLUMN IF NOT EXISTS step_id integer NOT NULL DEFAULT 0`,
	`ALTER TABLE inspection_record ADD COLUMN IF NOT EXISTS step_name varchar(128) NOT NULL DEFAULT ''`,
	`ALTER TABLE inspection_record ADD COLUMN IF NOT EXISTS measured_value varchar(64) NOT NULL DEFAULT ''`,
}

// EnsureSchema ent 建表 + 幂等列补丁（只增，不删数据）
func EnsureSchema(client *ent.Client, sqlDB *sql.DB) error {
	if err := client.Schema.Create(context.Background()); err != nil {
		return err
	}
	for _, q := range schemaPatchSQL {
		if _, err := sqlDB.Exec(q); err != nil {
			return fmt.Errorf("schema patch failed: %v", err)
		}
	}
	return nil
}

// AutoMigrate 执行 ent 自动建表/补齐（只增，不删数据）。供命令行 migrate 使用。

func AutoMigrate(c DatabaseConf) error {
	client, sqlDB := MustNewDB(c)
	defer sqlDB.Close()
	return EnsureSchema(client, sqlDB)
}

// DropAllTables 清空 public schema 下所有业务表（破坏性，仅供命令行 reset-all 使用）。
func DropAllTables(c DatabaseConf) error {
	sqlDB, err := NewDB(c)
	if err != nil {
		return err
	}
	defer sqlDB.Close()

	rows, err := sqlDB.Query(`SELECT tablename FROM pg_tables WHERE schemaname='public'`)
	if err != nil {
		return err
	}
	var names []string
	for rows.Next() {
		var t string
		if err := rows.Scan(&t); err != nil {
			rows.Close()
			return err
		}
		names = append(names, t)
	}
	rows.Close()

	for _, t := range names {
		if _, err := sqlDB.Exec(`DROP TABLE IF EXISTS "` + t + `" CASCADE`); err != nil {
			return err
		}
	}
	return nil
}