SafeLine/management/webserver/pkg/database/postgres.go

package database

import (
	"gorm.io/driver/postgres"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
	"gorm.io/gorm/schema"

	"chaitin.cn/dev/go/errors"
	"chaitin.cn/patronus/safeline-2/management/webserver/pkg/config"
)

type PostgresDB struct {
	*gorm.DB
}

var (
	// db is not designed to be used directly, use method `GetDB` instead.
	db PostgresDB
)

func GetDB() *PostgresDB {
	return &db
}

func InitDB() error {
	dbConfig := config.GlobalConfig.DB

	if dbConfig.URL == "" {
		return errors.New("empty database url")
	}
	// URL also works, see https://godoc.org/github.com/lib/pq
	var (
		gormDB *gorm.DB
		err    error
	)

	gormConfig := &gorm.Config{
		NamingStrategy: schema.NamingStrategy{
			TablePrefix:   "mgt_", // table name prefix, table for `User` would be `t_users`
			SingularTable: true,   // use singular table name, table for `User` would be `user` with this option enabled
		},
		AllowGlobalUpdate: false,
		Logger:            logger.Default.LogMode(logger.Silent),
	}
	if dbConfig.LogSQL {
		gormConfig.Logger = logger.Default.LogMode(logger.Info)
	}

	gormDB, err = gorm.Open(postgres.Open(dbConfig.URL), gormConfig)

	if err != nil {
		return err
	}

	db.DB = gormDB

	return nil
}

func (db *PostgresDB) SetDB(gormDB *gorm.DB) {
	db.DB = gormDB
}