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 ''`, // 上生产线前剩余项(H 质量检验 / I 产品物料清单 / J 步骤附件 / M 工单字段·绩效) `ALTER TABLE inspection_record ADD COLUMN IF NOT EXISTS report_no varchar(64) NOT NULL DEFAULT ''`, `ALTER TABLE inspection_record ADD COLUMN IF NOT EXISTS material_code varchar(64) NOT NULL DEFAULT ''`, `ALTER TABLE inspection_record ADD COLUMN IF NOT EXISTS material_name varchar(128) NOT NULL DEFAULT ''`, `ALTER TABLE inspection_record ADD COLUMN IF NOT EXISTS manufacturer varchar(128) NOT NULL DEFAULT ''`, `ALTER TABLE inspection_record ADD COLUMN IF NOT EXISTS inspection_no varchar(64) NOT NULL DEFAULT ''`, `ALTER TABLE inspection_record ADD COLUMN IF NOT EXISTS attachment_ids varchar(512) NOT NULL DEFAULT ''`, `ALTER TABLE inspection_record ADD COLUMN IF NOT EXISTS disposal_type varchar(20) NOT NULL DEFAULT ''`, `ALTER TABLE inspection_record ADD COLUMN IF NOT EXISTS disposal_remark varchar(500) NOT NULL DEFAULT ''`, // work_order_bom 表名注解为 work_order_bom `ALTER TABLE work_order_bom ADD COLUMN IF NOT EXISTS related_standard varchar(255) NOT NULL DEFAULT ''`, `ALTER TABLE process_step ADD COLUMN IF NOT EXISTS attachment jsonb NOT NULL DEFAULT '[]'::jsonb`, `ALTER TABLE work_order ADD COLUMN IF NOT EXISTS created_by varchar(64) NOT NULL DEFAULT ''`, `ALTER TABLE work_order ADD COLUMN IF NOT EXISTS due_date timestamp`, `ALTER TABLE workpiece_process ADD COLUMN IF NOT EXISTS duration_sec integer NOT NULL DEFAULT 0`, } // 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 }