SunYF
cbd55338f1
add a powershell script that supports starting MES, WMS, WMSClient, Workstation and Dashboard projects with one command, includes build and run logic for backend services and npm dev for dashboard
98 lines
2.4 KiB
Go
98 lines
2.4 KiB
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
|
|
}
|
|
|
|
// AutoMigrate 执行 ent 自动建表/补齐（只增，不删数据）。供命令行 migrate 使用。
|
|
|
|
func AutoMigrate(c DatabaseConf) error {
|
|
client, sqlDB := MustNewDB(c)
|
|
defer sqlDB.Close()
|
|
return client.Schema.Create(context.Background())
|
|
}
|
|
|
|
// 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
|
|
}
|