ds-sdk/db_operations/db.go

package db_operations

import (
	"errors"
	"fmt"
	"gorm.io/driver/postgres"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
)

const (
	databaseTypePostgres = "postgres"
)

type DBConfig struct {
	Type               string `mapstructure:"type"`
	UserName           string `mapstructure:"user_name"`
	Password           string `mapstructure:"password"`
	Address            string `mapstructure:"address"`
	Port               string `mapstructure:"port"`
	Database           string `mapstructure:"database"`
	MaxConnections     int    `mapstructure:"max_connections"`
	MaxIdleConnections int    `mapstructure:"max_idle_connections"`
}

func newGormDB(dbConfig *DBConfig) (*gorm.DB, error) {
	if dbConfig == nil {
		return nil, errors.New("没有传递数据库配置")
	}

	var gormDB *gorm.DB

	switch dbConfig.Type {
	case databaseTypePostgres:
		innerGormDB, err := newPostgresSQLOperationsItem(dbConfig)
		if err != nil {
			return nil, err
		}

		gormDB = innerGormDB
	default:
		innerGormDB, err := newPostgresSQLOperationsItem(dbConfig)
		if err != nil {
			return nil, err
		}

		gormDB = innerGormDB
	}

	return gormDB, nil
}

func newPostgresSQLOperationsItem(dbConfig *DBConfig) (*gorm.DB, error) {
	dsn := "host=%s port=%s user=%s password=%s dbname=%s sslmode=disable TimeZone=Asia/Shanghai"
	connStr := fmt.Sprintf(dsn, dbConfig.Address, dbConfig.Port, dbConfig.UserName, dbConfig.Password, dbConfig.Database)
	return gorm.Open(postgres.Open(connStr), &gorm.Config{
		Logger:      logger.Default.LogMode(logger.Info),
		PrepareStmt: true,
	})
}

func destroyGormDB(gormDB *gorm.DB) error {
	if gormDB == nil {
		return nil
	}

	db, err := gormDB.DB()
	if err != nil {
		return err
	}

	return db.Close()
}