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
}

// 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
}